[Feature] TorchRL2Gym conversion (pytorch#1795)

ai4co · Jan 19, 2024 · c3ffb5a · c3ffb5a
1 parent 57139bd
commit c3ffb5a
Show file tree

Hide file tree

Showing 23 changed files with 1,964 additions and 141 deletions.
diff --git a/docs/source/reference/envs.rst b/docs/source/reference/envs.rst
@@ -604,6 +604,7 @@ to be able to create this other composition:
     RewardScaling
     RewardSum
     Reward2GoTransform
+    RemoveEmptySpecs
     SelectTransform
     SignTransform
     SqueezeTransform

diff --git a/test/_utils_internal.py b/test/_utils_internal.py
@@ -63,7 +63,7 @@ def _set_gym_environments():  # noqa: F811
     PONG_VERSIONED = "ALE/Pong-v5"
 
 
-@implement_for("gymnasium", "0.27.0", None)
+@implement_for("gymnasium")
 def _set_gym_environments():  # noqa: F811
     global CARTPOLE_VERSIONED, HALFCHEETAH_VERSIONED, PENDULUM_VERSIONED, PONG_VERSIONED
 

diff --git a/test/mocking_classes.py b/test/mocking_classes.py
@@ -515,7 +515,7 @@ def _step(
             assert (a.sum(-1) == 1).all()
 
         obs = self._get_in_obs(tensordict.get(self._out_key)) + a / self.maxstep
-        tensordict = tensordict.empty()  # empty tensordict
+        tensordict = tensordict.empty()
 
         tensordict.set(self.out_key, self._get_out_obs(obs))
         tensordict.set(self._out_key, self._get_out_obs(obs))
@@ -603,6 +603,7 @@ def _reset(self, tensordict: TensorDictBase) -> TensorDictBase:
         # state = torch.zeros(self.size) + self.counter
         if tensordict is None:
             tensordict = TensorDict({}, self.batch_size, device=self.device)
+
         tensordict = tensordict.empty()
         tensordict.update(self.observation_spec.rand())
         # tensordict.set("next_" + self.out_key, self._get_out_obs(state))
@@ -622,6 +623,7 @@ def _step(
         a = tensordict.get("action")
 
         obs = self._obs_step(self._get_in_obs(tensordict.get(self._out_key)), a)
+
         tensordict = tensordict.empty()  # empty tensordict
 
         tensordict.set(self.out_key, self._get_out_obs(obs))