RandomActionPolicy: fix unpacking of act, state

thu-ml · Aug 14, 2024 · 02ab06a · 02ab06a
1 parent c699e1d
commit 02ab06a
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tianshou/policy/base.py b/tianshou/policy/base.py
@@ -711,8 +711,9 @@ def forward(
         batch: ObsBatchProtocol,
         state: dict | BatchProtocol | np.ndarray | None = None,
         **kwargs: Any,
-    ) -> ActBatchProtocol:
-        return cast(ActBatchProtocol, Batch(act=self.actor(batch.obs)))
+    ) -> ActStateBatchProtocol:
+        act, next_state = self.actor(batch.obs, state)
+        return cast(ActStateBatchProtocol, Batch(act=act, state=next_state))
 
     def learn(self, batch: RolloutBatchProtocol, *args: Any, **kwargs: Any) -> TrainingStats:
         return TrainingStats()