feedback from vmoens

pytorch · vmoens · Nov 19, 2022 · Nov 2, 2022 · Nov 7, 2022 · Nov 7, 2022
commit d569512198d5b98acbb1c0bcf9b61abd6b48c16f
diff --git a/test/test_libs.py b/test/test_libs.py
@@ -322,8 +322,9 @@ def test_habitat(self, envname):
 
 
 @pytest.mark.skipif(not _has_jumanji, reason="jumanji not installed")
+@pytest.mark.parametrize("envname", ["Snake-6x6-v0", "TSP50-v0"])
 class TestJumanji:
-    @pytest.mark.parametrize("envname", ["Snake-6x6-v0"])
+
     def test_jumanji_seeding(self, envname):
         final_seed = []
         tdreset = []
@@ -341,17 +342,28 @@ def test_jumanji_seeding(self, envname):
         assert_allclose_td(*tdreset)
         assert_allclose_td(*tdrollout)
 
-    @pytest.mark.parametrize("batch_size", [(), (2,), (2, 3)])
-    @pytest.mark.parametrize("envname", ["Snake-6x6-v0"])
+    @pytest.mark.parametrize("batch_size", [(), (5,), (5, 4)])
     def test_jumanji_batch_size(self, envname, batch_size):
         env = JumanjiEnv(envname, batch_size=batch_size)
+        env.set_seed(0)
         tdreset = env.reset()
         tdrollout = env.rollout(max_steps=50)
         env.close()
         del env
         assert tdreset.batch_size == batch_size
         assert tdrollout.batch_size[:-1] == batch_size
 
+    @pytest.mark.parametrize("batch_size", [(), (5,), (5, 4)])
+    def test_jumanji_spec_rollout(self, envname, batch_size):
+        env = JumanjiEnv(envname, batch_size=batch_size)
+        env.set_seed(0)
+        tdrollout = env.rollout(max_steps=50)
+        fake_td = (
+            env.fake_tensordict().unsqueeze(-1).expand(*tdrollout.shape).contiguous()
+        )
+        tdrollout.zero_()
+        assert (tdrollout == fake_td).all()
+
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()

diff --git a/torchrl/envs/libs/jumanji.py b/torchrl/envs/libs/jumanji.py
@@ -74,7 +74,7 @@ def _jumanji_to_torchrl_spec_transform(
             )
         return CompositeSpec(**new_spec)
     else:
-        raise NotImplementedError(type(spec))
+        raise TypeError(f"Unsupported spec type {type(spec)}")
 
 
 def _jumanji_to_torchrl_obs_spec_transform(
@@ -91,7 +91,7 @@ def _jumanji_to_torchrl_obs_spec_transform(
     elif isinstance(spec, jumanji.specs.Spec):
         return CompositeSpec(**{f"next_{k}": v for k, v in new_spec.items()})
     else:
-        raise NotImplementedError(type(spec))
+        raise TypeError(f"Unsupported spec type {type(spec)}")
 
 
 def _jumanji_to_torchrl_state_spec_transform(
@@ -119,7 +119,7 @@ def _jumanji_to_torchrl_state_spec_transform(
             }
         )
     else:
-        raise NotImplementedError(type(state))
+        raise TypeError(f"Unsupported state type {type(state)}")
 
 
 def _jumanji_to_torchrl_input_spec_transform(
@@ -130,13 +130,14 @@ def _jumanji_to_torchrl_input_spec_transform(
     categorical_action_encoding: bool = True,
 ) -> TensorSpec:
     state_dict = _jumanji_to_torchrl_data_transform(state, device=device)
-    input_spec = _jumanji_to_torchrl_state_spec_transform(
-        state_dict, dtype, device, categorical_action_encoding
-    )
-    action_spec = _jumanji_to_torchrl_spec_transform(
-        action_spec, dtype, device, categorical_action_encoding
+    input_spec = CompositeSpec(
+        state=_jumanji_to_torchrl_state_spec_transform(
+            state_dict, dtype, device, categorical_action_encoding
+        ),
+        action=_jumanji_to_torchrl_spec_transform(
+            action_spec, dtype, device, categorical_action_encoding
+        )
     )
-    input_spec["action"] = action_spec
     return input_spec
 
 
@@ -151,17 +152,18 @@ def _jumanji_to_torchrl_data_transform(val, device):
         if val.dtype == np.uint64:
             val = val.astype(np.int64)
         return torch.tensor(val, device=device)
-    if isinstance(val, tuple) and hasattr(val, "_fields"):  # named tuples
+    elif isinstance(val, tuple) and hasattr(val, "_fields"):  # named tuples
         return {
             k: _jumanji_to_torchrl_data_transform(v, device=device)
             for k, v in zip(val._fields, val)
         }
-    if hasattr(val, "__dict__"):
+    elif hasattr(val, "__dict__"):
         return {
             k: _jumanji_to_torchrl_data_transform(v, device=device)
             for k, v in val.__dict__.items()
         }
-    raise TypeError(f"Unsupported data type {type(val)}")
+    else:
+        raise TypeError(f"Unsupported data type {type(val)}")
 
 
 def _torchrl_to_jumanji_state_transform(tensordict: TensorDict, env):
@@ -287,9 +289,8 @@ def _build_env(
 
     def _make_specs(self, env: "jumanji.env.Environment") -> None:  # noqa: F821
         # generate a sample state object to build state spec from.
-        seed = int.from_bytes(np.random.bytes(7), byteorder="big", signed=False)
-        self.set_seed(seed)
-        state, _ = env.reset(self.key)
+        key = jax.random.PRNGKey(0)
+        state, _ = env.reset(key)
 
         self._input_spec = _jumanji_to_torchrl_input_spec_transform(
             env.action_spec(), state, device=self.device
@@ -308,14 +309,23 @@ def _check_kwargs(self, kwargs: Dict):
         if not isinstance(env, (jumanji.env.Environment,)):
             raise TypeError("env is not of type 'jumanji.env.Environment'.")
 
-    def _init_env(self) -> Optional[int]:
+    def _init_env(self):
         pass
 
     def _set_seed(self, seed):
         if seed is None:
             raise Exception("Jumanji requires an integer seed.")
         self.key = jax.random.PRNGKey(seed)
 
+    def read_state(self, state):
+        state = _jumanji_to_torchrl_data_transform(state, device=self.device)
+        state = self.input_spec["state"].encode(state)
+        return state
+
+    def read_obs(self, obs):
+        obs = _jumanji_to_torchrl_data_transform(obs, device=self.device)
+        return super().read_obs(obs)
+
     def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
 
         state = _torchrl_to_jumanji_state_transform(tensordict.get("state"), self._env)
@@ -328,16 +338,11 @@ def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
         state = self._reshape(state)
         timestep = self._reshape(timestep)
 
-        state_dict = _jumanji_to_torchrl_data_transform(state, device=self.device)
-        obs_dict = self.read_obs(
-            _jumanji_to_torchrl_data_transform(timestep.observation, device=self.device)
-        )
-        reward = self.read_reward(
-            reward,
-            _jumanji_to_torchrl_data_transform(timestep.reward, device=self.device),
-        )
-        done = _jumanji_to_torchrl_data_transform(
-            timestep.step_type == self.lib.types.StepType.LAST, device=self.device
+        state_dict = self.read_state(state)
+        obs_dict = self.read_obs(timestep.observation)
+        reward = self.read_reward(reward, np.asarray(timestep.reward))
+        done = torch.tensor(
+            np.asarray(timestep.step_type == self.lib.types.StepType.LAST)
         )
 
         self._is_done = done
@@ -349,7 +354,7 @@ def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
         )
         tensordict_out.set("reward", reward)
         tensordict_out.set("done", done)
-        tensordict_out.set("state", state_dict)
+        tensordict_out["state"] = state_dict
 
         return tensordict_out
 
@@ -362,11 +367,9 @@ def _reset(
         state = self._reshape(state)
         timestep = self._reshape(timestep)
 
-        state_dict = _jumanji_to_torchrl_data_transform(state, device=self.device)
-        obs_dict = self.read_obs(
-            _jumanji_to_torchrl_data_transform(timestep.observation, device=self.device)
-        )
-        done = torch.zeros(self.batch_size, dtype=torch.bool, device=self.device)
+        state_dict = self.read_state(state)
+        obs_dict = self.read_obs(timestep.observation)
+        done = torch.zeros(self.batch_size, dtype=torch.bool)
 
         self._is_done = done
 
@@ -376,7 +379,7 @@ def _reset(
             device=self.device,
         )
         tensordict_out.set("done", done)
-        tensordict_out.set("state", state_dict)
+        tensordict_out["state"] = state_dict
 
         return tensordict_out
 
@@ -397,7 +400,6 @@ class JumanjiEnv(JumanjiWrapper):
         >>> td = env.rand_step()
         >>> print(td)
         >>> print(env.available_envs)
-
     """
 
     def __init__(self, env_name, **kwargs):