[Feature] InitTracker transform (#962)

Author: vmoens

docs/source/reference/envs.rst

@@ -49,7 +49,8 @@ With these, the following methods are implemented:
 
 - :meth:`env.reset`: a reset method that may (but not necessarily requires to) take
   a :class:`tensordict.TensorDict` input. It return the first tensordict of a rollout, usually
-  containing a :obj:`"done"` state and a set of observations.
+  containing a :obj:`"done"` state and a set of observations. If not present,
+  a `"reward"` key will be instantiated with 0s and the appropriate shape.
 - :meth:`env.step`: a step method that takes a :class:`tensordict.TensorDict` input
   containing an input action as well as other inputs (for model-based or stateless
   environments, for instance).
@@ -88,6 +89,21 @@ function.
   TorchRL's collectors and rollout methods will be looking for one of these
   keys when assessing if the env should be reset.
 
+.. note::
+
+  The `torchrl.collectors.utils.split_trajectories` function can be used to
+  slice adjacent trajectories. It relies on a ``"traj_ids"`` entry in the
+  input tensordict, or to the junction of ``"done"`` and ``"truncated"`` key
+  if the ``"traj_ids"`` is missing.
+
+
+.. note::
+
+  In some contexts, it can be useful to mark the first step of a trajectory.
+  TorchRL provides such functionality through the :class:`torchrl.envs.InitTracker`
+  transform.
+
+
 Our environment `tutorial <https://pytorch.org/rl/tutorials/pendulum.html>`_
 provides more information on how to design a custom environment from scratch.
 
@@ -309,6 +325,7 @@ to be able to create this other composition:
     FrameSkipTransform
     GrayScale
     gSDENoise
+    InitTracker
     NoopResetEnv
     ObservationNorm
     ObservationTransform

test/mocking_classes.py

@@ -905,7 +905,8 @@ def __init__(self, max_steps: int = 5, start_val: int = 0, **kwargs):
                 (
                     *self.batch_size,
                     1,
-                )
+                ),
+                dtype=torch.int32,
             ),
             shape=self.batch_size,
         )
@@ -915,6 +916,14 @@ def __init__(self, max_steps: int = 5, start_val: int = 0, **kwargs):
                 1,
             )
         )
+        self.done_spec = DiscreteTensorSpec(
+            2,
+            dtype=torch.bool,
+            shape=(
+                *self.batch_size,
+                1,
+            ),
+        )
         self.input_spec = CompositeSpec(
             action=BinaryDiscreteTensorSpec(n=1, shape=[*self.batch_size, 1]),
             shape=self.batch_size,
@@ -978,19 +987,19 @@ def __init__(
         if max_steps is None:
             max_steps = torch.tensor(5)
         if start_val is None:
-            start_val = torch.zeros(())
+            start_val = torch.zeros((), dtype=torch.int32)
         if not max_steps.shape == self.batch_size:
             raise RuntimeError("batch_size and max_steps shape must match.")
 
         self.max_steps = max_steps
-        self.start_val = start_val
 
         self.observation_spec = CompositeSpec(
             observation=UnboundedContinuousTensorSpec(
                 (
                     *self.batch_size,
                     1,
-                )
+                ),
+                dtype=torch.int32,
             ),
             shape=self.batch_size,
         )
@@ -1000,6 +1009,14 @@ def __init__(
                 1,
             )
         )
+        self.done_spec = DiscreteTensorSpec(
+            2,
+            dtype=torch.bool,
+            shape=(
+                *self.batch_size,
+                1,
+            ),
+        )
         self.input_spec = CompositeSpec(
             action=BinaryDiscreteTensorSpec(n=1, shape=[*self.batch_size, 1]),
             shape=self.batch_size,
@@ -1008,20 +1025,24 @@ def __init__(
         self.count = torch.zeros(
             (*self.batch_size, 1), device=self.device, dtype=torch.int
         )
+        if start_val.numel() == self.batch_size.numel():
+            self.start_val = start_val.view(*self.batch_size, 1)
+        elif start_val.numel() <= 1:
+            self.start_val = start_val.expand_as(self.count)
 
     def _set_seed(self, seed: Optional[int]):
         torch.manual_seed(seed)
 
     def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
         if tensordict is not None and "_reset" in tensordict.keys():
             _reset = tensordict.get("_reset")
-            self.count[_reset] = self.start_val[_reset].unsqueeze(-1)
+            self.count[_reset] = self.start_val[_reset].view_as(self.count[_reset])
         else:
-            self.count[:] = self.start_val.unsqueeze(-1)
+            self.count[:] = self.start_val.view_as(self.count)
         return TensorDict(
             source={
                 "observation": self.count.clone(),
-                "done": self.count > self.max_steps.unsqueeze(-1),
+                "done": self.count > self.max_steps.view_as(self.count),
             },
             batch_size=self.batch_size,
             device=self.device,
@@ -1032,7 +1053,7 @@ def _step(
         tensordict: TensorDictBase,
     ) -> TensorDictBase:
         action = tensordict.get("action")
-        self.count += action.to(torch.int).unsqueeze(-1)
+        self.count += action.to(torch.int).view_as(self.count)
         tensordict = TensorDict(
             source={
                 "observation": self.count.clone(),

test/test_collector.py

@@ -1223,17 +1223,17 @@ def test_initial_obs_consistency(env_class, seed=1):
     if env_class == CountingEnv:
         arange_0 = start_val + torch.arange(max_steps - 3)
         arange = start_val + torch.arange(2)
-        expected = torch.cat([arange_0, arange_0, arange]).float()
+        expected = torch.cat([arange_0, arange_0, arange])
     else:
         # the first env has a shorter horizon than the second
         arange_0 = start_val + torch.arange(max_steps - 3 - 1)
         arange = start_val + torch.arange(start_val)
-        expected_0 = torch.cat([arange_0, arange_0, arange]).float()
+        expected_0 = torch.cat([arange_0, arange_0, arange])
         arange_0 = start_val + torch.arange(max_steps - 3)
         arange = start_val + torch.arange(2)
-        expected_1 = torch.cat([arange_0, arange_0, arange]).float()
+        expected_1 = torch.cat([arange_0, arange_0, arange])
         expected = torch.stack([expected_0, expected_1])
-    assert torch.allclose(obs, expected)
+    assert torch.allclose(obs, expected.to(obs.dtype))
 
 
 def weight_reset(m):

test/test_transforms.py

@@ -21,6 +21,7 @@
 )
 from mocking_classes import (
     ContinuousActionVecMockEnv,
+    CountingBatchedEnv,
     DiscreteActionConvMockEnvNumpy,
     MockBatchedLockedEnv,
     MockBatchedUnLockedEnv,
@@ -75,7 +76,7 @@
 from torchrl.envs.libs.gym import _has_gym, GymEnv
 from torchrl.envs.transforms import VecNorm
 from torchrl.envs.transforms.r3m import _R3MNet
-from torchrl.envs.transforms.transforms import _has_tv
+from torchrl.envs.transforms.transforms import _has_tv, InitTracker
 from torchrl.envs.transforms.vip import _VIPNet, VIPRewardTransform
 from torchrl.envs.utils import check_env_specs, step_mdp
 
@@ -6314,6 +6315,99 @@ def test_crop_mask(self, mask_key):
         assert tensordict_crop[mask_key].all()
 
 
+class TestInitTracker(TransformBase):
+    def test_single_trans_env_check(self):
+        env = CountingBatchedEnv(max_steps=torch.tensor([4, 5]), batch_size=[2])
+        env = TransformedEnv(env, InitTracker())
+        check_env_specs(env)
+
+    def test_serial_trans_env_check(self):
+        def make_env():
+            env = CountingBatchedEnv(max_steps=torch.tensor([4, 5]), batch_size=[2])
+            env = TransformedEnv(env, InitTracker())
+            return env
+
+        env = SerialEnv(2, make_env)
+        check_env_specs(env)
+
+    def test_parallel_trans_env_check(self):
+        def make_env():
+            env = CountingBatchedEnv(max_steps=torch.tensor([4, 5]), batch_size=[2])
+            env = TransformedEnv(env, InitTracker())
+            return env
+
+        env = ParallelEnv(2, make_env)
+        check_env_specs(env)
+
+    def test_trans_serial_env_check(self):
+        def make_env():
+            env = CountingBatchedEnv(max_steps=torch.tensor([4, 5]), batch_size=[2])
+            return env
+
+        env = SerialEnv(2, make_env)
+        env = TransformedEnv(env, InitTracker())
+        check_env_specs(env)
+
+    def test_trans_parallel_env_check(self):
+        def make_env():
+            env = CountingBatchedEnv(max_steps=torch.tensor([4, 5]), batch_size=[2])
+            return env
+
+        env = ParallelEnv(2, make_env)
+        env = TransformedEnv(env, InitTracker())
+        check_env_specs(env)
+
+    def test_transform_no_env(self):
+        with pytest.raises(
+            NotImplementedError, match="InitTracker cannot be executed without a parent"
+        ):
+            InitTracker()(None)
+
+    def test_transform_compose(self):
+        with pytest.raises(
+            NotImplementedError, match="InitTracker cannot be executed without a parent"
+        ):
+            Compose(InitTracker())(None)
+
+    def test_transform_env(self):
+        policy = lambda tensordict: tensordict.set(
+            "action", torch.ones(tensordict.shape, dtype=torch.int32)
+        )
+        env = CountingBatchedEnv(max_steps=torch.tensor([3, 4]), batch_size=[2])
+        env = TransformedEnv(env, InitTracker())
+        r = env.rollout(100, policy, break_when_any_done=False)
+        assert (r["is_init"].sum(1) == torch.tensor([25, 20])).all()
+
+    def test_transform_model(self):
+        with pytest.raises(
+            NotImplementedError, match="InitTracker cannot be executed without a parent"
+        ):
+            td = TensorDict({}, [])
+            chain = nn.Sequential(InitTracker())
+            chain(td)
+
+    def test_transform_rb(self):
+        batch = [1]
+        device = "cpu"
+        rb = ReplayBuffer(LazyTensorStorage(20))
+        rb.append_transform(InitTracker())
+        reward = torch.randn(*batch, 1, device=device)
+        misc = torch.randn(*batch, 1, device=device)
+        td = TensorDict(
+            {"misc": misc, "reward": reward},
+            batch,
+            device=device,
+        )
+        rb.extend(td)
+        with pytest.raises(
+            NotImplementedError, match="InitTracker cannot be executed without a parent"
+        ):
+            _ = rb.sample(20)
+
+    def test_transform_inverse(self):
+        raise pytest.skip("No inverse for InitTracker")
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/envs/__init__.py

@@ -21,6 +21,7 @@
     FrameSkipTransform,
     GrayScale,
     gSDENoise,
+    InitTracker,
     NoopResetEnv,
     ObservationNorm,
     ObservationTransform,

torchrl/envs/common.py

@@ -708,7 +708,7 @@ def policy(td):
                 exclude_action=False,
             )
             if not break_when_any_done and done.any():
-                _reset = done.squeeze(-1)
+                _reset = done.view(tensordict.shape)
                 tensordict.set("_reset", _reset)
                 self.reset(tensordict)
 

torchrl/envs/transforms/__init__.py

@@ -18,6 +18,7 @@
     FrameSkipTransform,
     GrayScale,
     gSDENoise,
+    InitTracker,
     NoopResetEnv,
     ObservationNorm,
     ObservationTransform,

torchrl/envs/transforms/transforms.py

@@ -23,6 +23,7 @@
     CompositeSpec,
     ContinuousBox,
     DEVICE_TYPING,
+    DiscreteTensorSpec,
     OneHotDiscreteTensorSpec,
     TensorSpec,
     UnboundedContinuousTensorSpec,
@@ -56,6 +57,8 @@ def interpolation_fn(interpolation):  # noqa: D103
 IMAGE_KEYS = ["pixels"]
 _MAX_NOOPS_TRIALS = 10
 
+FORWARD_NOT_IMPLEMENTED = "class {} cannot be executed without a parent" "environment."
+
 
 def _apply_to_composite(function):
     def new_fun(self, observation_spec):
@@ -2960,6 +2963,11 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
         observation_spec.update(episode_specs)
         return observation_spec
 
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        raise NotImplementedError(
+            FORWARD_NOT_IMPLEMENTED.format(self.__class__.__name__)
+        )
+
 
 class StepCounter(Transform):
     """Counts the steps from a reset and sets the done state to True after a certain number of steps.
@@ -3370,3 +3378,63 @@ def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
         arange = arange.view(arange_shape)
         idx = idx_0 + arange
         return tensordict.gather(dim=self.sample_dim, index=idx)
+
+
+class InitTracker(Transform):
+    """Reset tracker.
+
+    This transform populates the step/reset tensordict with a reset tracker entry
+    that is set to ``True`` whenever :meth:`~.reset` is called.
+
+    Args:
+         init_key (str, optional): the key to be used for the tracker entry.
+
+    Examples:
+        >>> from torchrl.envs.libs.gym import GymEnv
+        >>> env = TransformedEnv(GymEnv("Pendulum-v1"), InitTracker())
+        >>> td = env.reset()
+        >>> print(td["is_init"])
+        tensor(True)
+        >>> td = env.rand_step(td)
+        >>> print(td["next", "is_init"])
+        tensor(False)
+
+    """
+
+    def __init__(self, init_key: bool = "is_init"):
+        super().__init__(in_keys=[], out_keys=[init_key])
+
+    def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        if self.out_keys[0] not in tensordict.keys():
+            device = tensordict.device
+            if device is None:
+                device = torch.device("cpu")
+            tensordict.set(
+                self.out_keys[0],
+                torch.zeros(tensordict.shape, device=device, dtype=torch.bool),
+            )
+        return tensordict
+
+    def reset(self, tensordict: TensorDictBase) -> TensorDictBase:
+        device = tensordict.device
+        if device is None:
+            device = torch.device("cpu")
+        _reset = tensordict.get("_reset", None)
+        if _reset is None:
+            _reset = torch.ones(tensordict.shape, device=device, dtype=torch.bool)
+        tensordict.set(self.out_keys[0], _reset.clone())
+        return tensordict
+
+    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        observation_spec[self.out_keys[0]] = DiscreteTensorSpec(
+            2,
+            dtype=torch.bool,
+            device=self.parent.device,
+            shape=self.parent.batch_size,
+        )
+        return observation_spec
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        raise NotImplementedError(
+            FORWARD_NOT_IMPLEMENTED.format(self.__class__.__name__)
+        )