[Doc] Better doc for Transform class

ghstack-source-id: 16e563b
Pull Request resolved: #2797

Author: Vincent Moens

docs/source/reference/envs.rst

@@ -865,6 +865,8 @@ The inverse process is executed with the output tensordict, where the `in_keys`
 
    Rename transform logic
 
+.. note:: During a call to `inv`, the transforms are executed in reversed order (compared to the forward / step mode).
+
 Transforming Tensors and Specs
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -900,6 +902,74 @@ tensor that should not be generated when using :meth:`~torchrl.envs.EnvBase.rand
 environment. Instead, `"action_discrete"` should be generated, and its continuous counterpart obtained from the
 transform. Therefore, the user should see the `"action_discrete"` entry being exposed, but not `"action"`.
 
+Designing your own Transform
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To create a basic, custom transform, you need to subclass the `Transform` class and implement the
+:meth:`~torchrl.envs._apply_transform` method. Here's an example of a simple transform that adds 1 to the observation
+tensor:
+
+    >>> class AddOneToObs(Transform):
+    ...     """A transform that adds 1 to the observation tensor."""
+    ...
+    ...     def __init__(self):
+    ...         super().__init__(in_keys=["observation"], out_keys=["observation"])
+    ...
+    ...     def _apply_transform(self, obs: torch.Tensor) -> torch.Tensor:
+    ...         return obs + 1
+
+
+Tips for subclassing `Transform`
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+There are various ways of subclassing a transform. The things to take into considerations are:
+
+- Is the transform identical for each tensor / item being transformed? Use
+  :meth:`~torchrl.envs.Transform._apply_transform` and :meth:`~torchrl.envs.Transform._inv_apply_transform`.
+- The transform needs access to the input data to env.step as well as output? Rewrite
+  :meth:`~torchrl.envs.Transform._step`.
+  Otherwise, rewrite :meth:`~torchrl.envs.Transform._call` (or :meth:`~torchrl.envs.Transform._inv_call`).
+- Is the transform to be used within a replay buffer? Overwrite :meth:`~torchrl.envs.Transform.forward`,
+  :meth:`~torchrl.envs.Transform.inv`, :meth:`~torchrl.envs.Transform._apply_transform` or
+  :meth:`~torchrl.envs.Transform._inv_apply_transform`.
+- Within a transform, you can access (and make calls to) the parent environment using
+  :attr:`~torchrl.envs.Transform.parent` (the base env + all transforms till this one) or
+  :meth:`~torchrl.envs.Transform.container` (The object that encapsulates the transform).
+- Don't forget to edits the specs if needed: top level: :meth:`~torchrl.envs.Transform.transform_output_spec`,
+  :meth:`~torchrl.envs.Transform.transform_input_spec`.
+  Leaf level: :meth:`~torchrl.envs.Transform.transform_observation_spec`,
+  :meth:`~torchrl.envs.Transform.transform_action_spec`, :meth:`~torchrl.envs.Transform.transform_state_spec`,
+  :meth:`~torchrl.envs.Transform.transform_reward_spec` and
+  :meth:`~torchrl.envs.Transform.transform_reward_spec`.
+
+For practical examples, see the methods listed above.
+
+You can use a transform in an environment by passing it to the TransformedEnv constructor:
+
+    >>> env = TransformedEnv(GymEnv("Pendulum-v1"), AddOneToObs())
+
+You can compose multiple transforms together using the Compose class:
+
+    >>> transform = Compose(AddOneToObs(), RewardSum())
+    >>> env = TransformedEnv(GymEnv("Pendulum-v1"), transform)
+
+Inverse Transforms
+^^^^^^^^^^^^^^^^^^
+
+Some transforms have an inverse transform that can be used to undo the transformation. For example, the AddOneToAction
+transform has an inverse transform that subtracts 1 from the action tensor:
+
+    >>> class AddOneToAction(Transform):
+    ...     """A transform that adds 1 to the action tensor."""
+    ...     def __init__(self):
+    ...         super().__init__(in_keys=[], out_keys=[], in_keys_inv=["action"], out_keys_inv=["action"])
+    ...     def _inv_apply_transform(self, action: torch.Tensor) -> torch.Tensor:
+    ...         return action + 1
+
+Using a Transform with a Replay Buffer
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+You can use a transform with a replay buffer by passing it to the ReplayBuffer constructor:
 
 Cloning transforms
 ~~~~~~~~~~~~~~~~~~

test/test_collector.py

@@ -3213,9 +3213,11 @@ def test_cudagraph_policy(self, collector_cls, cudagraph_policy):
 @pytest.mark.skipif(not _has_gym, reason="gym required for this test")
 class TestCollectorsNonTensor:
     class AddNontTensorData(Transform):
-        def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
-            tensordict["nt"] = f"a string! - {tensordict.get('step_count').item()}"
-            return tensordict
+        def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
+            next_tensordict[
+                "nt"
+            ] = f"a string! - {next_tensordict.get('step_count').item()}"
+            return next_tensordict
 
         def _reset(
             self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase

torchrl/envs/model_based/dreamer.py

@@ -80,8 +80,8 @@ class DreamerDecoder(Transform):
         >>> model_based_env_eval = model_based_env.append_transform(DreamerDecoder())
     """
 
-    def _call(self, tensordict):
-        return self.parent.base_env.obs_decoder(tensordict)
+    def _call(self, next_tensordict):
+        return self.parent.base_env.obs_decoder(next_tensordict)
 
     def _reset(self, tensordict, tensordict_reset):
         return self._call(tensordict_reset)

torchrl/envs/transforms/gym_transforms.py

@@ -138,8 +138,8 @@ def _get_lives(self):
             lives = torch.as_tensor([_lives() for _lives in lives])
         return lives
 
-    def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
-        return tensordict
+    def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
+        return next_tensordict
 
     def _step(self, tensordict, next_tensordict):
         parent = self.parent

torchrl/envs/transforms/r3m.py

@@ -70,12 +70,12 @@ def __init__(self, in_keys, out_keys, model_name, del_keys: bool = True):
         self.del_keys = del_keys
 
     @set_lazy_legacy(False)
-    def _call(self, tensordict):
-        with tensordict.view(-1) as tensordict_view:
+    def _call(self, next_tensordict):
+        with next_tensordict.view(-1) as tensordict_view:
             super()._call(tensordict_view)
         if self.del_keys:
-            tensordict.exclude(*self.in_keys, inplace=True)
-        return tensordict
+            next_tensordict.exclude(*self.in_keys, inplace=True)
+        return next_tensordict
 
     forward = _call
 

torchrl/envs/transforms/rlhf.py

@@ -158,25 +158,25 @@ def _reset(
             tensordict_reset = self._call(tensordict_reset)
         return tensordict_reset
 
-    def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
+    def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
         # run the actor on the tensordict
-        action = tensordict.get("action", None)
+        action = next_tensordict.get("action", None)
         if action is None:
             # being called after reset or without action, skipping
             if self.out_keys[0] != ("reward",) and self.parent is not None:
-                tensordict.set(self.out_keys[0], self.parent.reward_spec.zero())
-            return tensordict
+                next_tensordict.set(self.out_keys[0], self.parent.reward_spec.zero())
+            return next_tensordict
         with self.frozen_params.to_module(self.functional_actor):
-            dist = self.functional_actor.get_dist(tensordict.clone(False))
+            dist = self.functional_actor.get_dist(next_tensordict.clone(False))
         # get the log_prob given the original model
         log_prob = dist.log_prob(action)
         reward_key = self.in_keys[0]
-        reward = tensordict.get("next").get(reward_key)
-        curr_log_prob = tensordict.get(self.sample_log_prob_key)
+        reward = next_tensordict.get("next").get(reward_key)
+        curr_log_prob = next_tensordict.get(self.sample_log_prob_key)
         # we use the unbiased consistent estimator of the KL: log_p(x) - log_q(x) when x ~ p(x)
         kl = (curr_log_prob - log_prob).view_as(reward)
-        tensordict.set(("next", *self.out_keys[0]), reward + self.coef * kl)
-        return tensordict
+        next_tensordict.set(("next", *self.out_keys[0]), reward + self.coef * kl)
+        return next_tensordict
 
     def _step(
         self, tensordict: TensorDictBase, next_tensordict: TensorDictBase