[Environment] Complete PettingZooWrapper state support (#2953)

Co-authored-by: Vincent Moens <[email protected]>
(cherry picked from commit d882ea2)

Author: jeguzzi

test/test_libs.py

@@ -3944,6 +3944,20 @@ def __call__(self, td):
                 td[-1]["next", "player", "reward"] == torch.tensor([[-1], [1]])
             ).all()
 
+    @pytest.mark.parametrize("task", ["simple_v3"])
+    def test_return_state(self, task):
+        env = PettingZooEnv(
+            task=task,
+            parallel=True,
+            seed=0,
+            use_mask=False,
+            return_state=True,
+        )
+        check_env_specs(env)
+        r = env.rollout(10)
+        assert (r["state"] != 0).any()
+        assert (r["next", "state"] != 0).any()
+
     @pytest.mark.parametrize(
         "task",
         [

torchrl/envs/libs/pettingzoo.py

@@ -584,6 +584,10 @@ def _reset(
                             value, device=self.device
                         )
 
+        if self.return_state:
+            state = torch.as_tensor(self.state(), device=self.device)
+            tensordict_out.set("state", state)
+
         return tensordict_out
 
     def _reset_aec(self, **kwargs) -> tuple[dict, dict]:
@@ -702,6 +706,11 @@ def _step(
         tensordict_out.set("done", done)
         tensordict_out.set("terminated", terminated)
         tensordict_out.set("truncated", truncated)
+
+        if self.return_state:
+            state = torch.as_tensor(self.state(), device=self.device)
+            tensordict_out.set("state", state)
+
         return tensordict_out
 
     def _aggregate_done(self, tensordict_out, use_any):