[Feature] Timer transform

ghstack-source-id: e42f2ae
Pull Request resolved: #2806

Author: Vincent Moens

docs/source/reference/envs.rst

@@ -1070,6 +1070,7 @@ to be able to create this other composition:
     TargetReturn
     TensorDictPrimer
     TimeMaxPool
+    Timer
     Tokenizer
     ToTensorImage
     TrajCounter

test/test_transforms.py

@@ -24,6 +24,7 @@
 import tensordict.tensordict
 import torch
 from tensordict import (
+    LazyStackedTensorDict,
     NonTensorData,
     NonTensorStack,
     TensorDict,
@@ -102,6 +103,7 @@
     TargetReturn,
     TensorDictPrimer,
     TimeMaxPool,
+    Timer,
     Tokenizer,
     ToTensorImage,
     TrajCounter,
@@ -13879,6 +13881,90 @@ def test_transform_inverse(self):
         return
 
 
+class TestTimer(TransformBase):
+    def test_single_trans_env_check(self):
+        env = TransformedEnv(ContinuousActionVecMockEnv(), Timer())
+        check_env_specs(env)
+        env.close()
+
+    def test_serial_trans_env_check(self):
+        env = SerialEnv(
+            2, lambda: TransformedEnv(ContinuousActionVecMockEnv(), Timer())
+        )
+        check_env_specs(env)
+        env.close()
+
+    def test_parallel_trans_env_check(self, maybe_fork_ParallelEnv):
+        env = maybe_fork_ParallelEnv(
+            2, lambda: TransformedEnv(ContinuousActionVecMockEnv(), Timer())
+        )
+        try:
+            check_env_specs(env)
+        finally:
+            try:
+                env.close()
+            except RuntimeError:
+                pass
+
+    def test_trans_serial_env_check(self):
+        env = TransformedEnv(
+            SerialEnv(2, lambda: ContinuousActionVecMockEnv()), Timer()
+        )
+        try:
+            check_env_specs(env)
+        finally:
+            try:
+                env.close()
+            except RuntimeError:
+                pass
+
+    def test_trans_parallel_env_check(self, maybe_fork_ParallelEnv):
+        env = TransformedEnv(
+            maybe_fork_ParallelEnv(2, lambda: ContinuousActionVecMockEnv()),
+            Timer(),
+        )
+        try:
+            check_env_specs(env)
+        finally:
+            try:
+                env.close()
+            except RuntimeError:
+                pass
+
+    def test_transform_no_env(self):
+        torch.manual_seed(0)
+        t = Timer()
+        with pytest.raises(NotImplementedError):
+            t(TensorDict())
+
+    def test_transform_compose(self):
+        torch.manual_seed(0)
+        t = Compose(Timer())
+        with pytest.raises(NotImplementedError):
+            t(TensorDict())
+
+    def test_transform_env(self):
+        env = TransformedEnv(ContinuousActionVecMockEnv(), Timer())
+        rollout = env.rollout(3)
+        # The stack must be contiguous
+        assert not isinstance(rollout, LazyStackedTensorDict)
+        assert (rollout["time_policy"] >= 0).all()
+        assert (rollout["time_step"] > 0).all()
+
+    def test_transform_model(self):
+        torch.manual_seed(0)
+        t = nn.Sequential(Timer())
+        with pytest.raises(NotImplementedError):
+            t(TensorDict())
+
+    def test_transform_rb(self):
+        # NotImplemented tested elsewhere
+        return
+
+    def test_transform_inverse(self):
+        raise pytest.skip("Tested elsewhere")
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/envs/__init__.py

@@ -97,6 +97,7 @@
     TargetReturn,
     TensorDictPrimer,
     TimeMaxPool,
+    Timer,
     Tokenizer,
     ToTensorImage,
     TrajCounter,

torchrl/envs/transforms/__init__.py

@@ -57,6 +57,7 @@
     TargetReturn,
     TensorDictPrimer,
     TimeMaxPool,
+    Timer,
     Tokenizer,
     ToTensorImage,
     TrajCounter,

torchrl/envs/transforms/transforms.py

@@ -9,6 +9,7 @@
 import hashlib
 import importlib.util
 import multiprocessing as mp
+import time
 import warnings
 import weakref
 from copy import copy
@@ -10823,3 +10824,112 @@ def _transform_observation_spec(
             )
         )
         return observation_spec
+
+
+class Timer(Transform):
+    """A transform that measures the time intervals between `inv` and `call` operations in an environment.
+
+    The `Timer` transform is used to track the time elapsed between the `inv` call and the `call`,
+    and between the `call` and the `inv` call. This is useful for performance monitoring and debugging
+    within an environment. The time is measured in seconds and stored as a tensor with the default
+    dtype from PyTorch. If the tensordict has a batch size (e.g., in batched environments), the time will be expended
+    to the size of the input tensordict.
+
+    Attributes:
+        out_keys: The keys of the output tensordict for the inverse transform. Defaults to
+            `out_keys = [f"{time_key}_step", f"{time_key}_policy"]`, where the first key represents
+            the time it takes to make a step in the environment, and the second key represents the
+            time it takes to execute the policy.
+        time_key: A prefix for the keys where the time intervals will be stored in the tensordict.
+            Defaults to `"time"`.
+
+    Examples:
+        >>> from torchrl.envs import Timer, GymEnv
+        >>>
+        >>> env = GymEnv("Pendulum-v1").append_transform(Timer())
+        >>> r = env.rollout(10)
+        >>> print("time for policy", r["time_policy"])
+        time for policy tensor([0.0000, 0.0882, 0.0004, 0.0002, 0.0002, 0.0002, 0.0002, 0.0002, 0.0002,
+                0.0002])
+        >>> print("time for step", r["time_step"])
+        time for step tensor([9.5797e-04, 1.6289e-03, 9.7990e-05, 8.0824e-05, 9.0837e-05, 7.6056e-05,
+                8.2016e-05, 7.6056e-05, 8.1062e-05, 7.7009e-05])
+    """
+
+    def __init__(self, out_keys: Sequence[NestedKey] = None, time_key: str = "time"):
+        if out_keys is None:
+            out_keys = [f"{time_key}_step", f"{time_key}_policy"]
+        elif len(out_keys) != 2:
+            raise TypeError(f"Expected two out_keys. Got out_keys={out_keys}.")
+        super().__init__([], out_keys)
+        self.time_key = time_key
+        self.last_inv_time = None
+        self.last_call_time = None
+
+    def _reset_env_preprocess(self, tensordict: TensorDictBase) -> TensorDictBase:
+        self.last_inv_time = time.time()
+        return tensordict
+
+    def _maybe_expand_and_set(self, key, time_elapsed, tensordict):
+        if isinstance(key, tuple):
+            parent_td = tensordict.get(key[:-1])
+            key = key[-1]
+        else:
+            parent_td = tensordict
+        batch_size = parent_td.batch_size
+        if batch_size:
+            # Get the parent shape
+            time_elapsed_expand = time_elapsed.expand(parent_td.batch_size)
+        else:
+            time_elapsed_expand = time_elapsed
+        parent_td.set(key, time_elapsed_expand)
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        current_time = time.time()
+        if self.last_inv_time is not None:
+            time_elapsed = torch.tensor(
+                current_time - self.last_inv_time, device=tensordict.device
+            )
+            self._maybe_expand_and_set(self.out_keys[0], time_elapsed, tensordict_reset)
+        self.last_call_time = current_time
+        # Placeholder
+        self._maybe_expand_and_set(self.out_keys[1], time_elapsed * 0, tensordict_reset)
+        return tensordict_reset
+
+    def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        current_time = time.time()
+        if self.last_call_time is not None:
+            time_elapsed = torch.tensor(
+                current_time - self.last_call_time, device=tensordict.device
+            )
+            self._maybe_expand_and_set(self.out_keys[1], time_elapsed, tensordict)
+        self.last_inv_time = current_time
+        return tensordict
+
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        current_time = time.time()
+        if self.last_inv_time is not None:
+            time_elapsed = torch.tensor(
+                current_time - self.last_inv_time, device=tensordict.device
+            )
+            self._maybe_expand_and_set(self.out_keys[0], time_elapsed, next_tensordict)
+        self.last_call_time = current_time
+        # presumbly no need to worry about batch size incongruencies here
+        next_tensordict.set(self.out_keys[1], tensordict.get(self.out_keys[1]))
+        return next_tensordict
+
+    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        observation_spec[self.out_keys[0]] = Unbounded(
+            shape=observation_spec.shape, device=observation_spec.device
+        )
+        observation_spec[self.out_keys[1]] = Unbounded(
+            shape=observation_spec.shape, device=observation_spec.device
+        )
+        return observation_spec
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        raise NotImplementedError(FORWARD_NOT_IMPLEMENTED)