Small typing improvement, related to https://github.com/thu-ml/tianshou/pull/915#discussion_r1329734222