Fix var-annotated by providing type for policy var

dantp-ai · Jan 11, 2024 · fce2482 · fce2482
1 parent fe3a0e0
commit fce2482
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/offline/test_cql.py b/test/offline/test_cql.py
@@ -12,6 +12,7 @@
 from tianshou.data import Collector, VectorReplayBuffer
 from tianshou.env import DummyVectorEnv
 from tianshou.policy import BasePolicy, CQLPolicy
+from tianshou.policy.cql import CQLTrainingStats
 from tianshou.trainer import OfflineTrainer
 from tianshou.utils import TensorboardLogger
 from tianshou.utils.net.common import Net
@@ -129,7 +130,7 @@ def test_cql(args: argparse.Namespace = get_args()) -> None:
         alpha_optim = torch.optim.Adam([log_alpha], lr=args.alpha_lr)
         args.alpha = (target_entropy, log_alpha, alpha_optim)
 
-    policy = CQLPolicy(
+    policy: CQLPolicy[CQLTrainingStats] = CQLPolicy(
         actor=actor,
         actor_optim=actor_optim,
         critic=critic,