Update

[ghstack-poisoned]

Author: vmoens

test/test_env.py

@@ -1692,6 +1692,34 @@ def test_parallel_env_device(
             env_serial.close(raise_if_closed=False)
             env0.close(raise_if_closed=False)
 
+    @pytest.mark.skipif(not _has_gym, reason="no gym")
+    @pytest.mark.parametrize("env_device", [None, "cpu"])
+    def test_parallel_env_device_vs_no_device(self, maybe_fork_ParallelEnv, env_device):
+        def make_env() -> GymEnv:
+            env = GymEnv(PENDULUM_VERSIONED(), device=env_device)
+            return env.append_transform(DoubleToFloat())
+
+        # Rollouts work with a regular env
+        parallel_env = maybe_fork_ParallelEnv(
+            num_workers=1, create_env_fn=make_env, device=None
+        )
+        parallel_env.reset()
+        parallel_env.set_seed(0)
+        torch.manual_seed(0)
+
+        parallel_rollout = parallel_env.rollout(max_steps=10)
+
+        # Rollout doesn't work with Parallelnv
+        parallel_env = maybe_fork_ParallelEnv(
+            num_workers=1, create_env_fn=make_env, device="cpu"
+        )
+        parallel_env.reset()
+        parallel_env.set_seed(0)
+        torch.manual_seed(0)
+
+        parallel_rollout_cpu = parallel_env.rollout(max_steps=10)
+        assert_allclose_td(parallel_rollout, parallel_rollout_cpu)
+
     @pytest.mark.skipif(not _has_gym, reason="no gym")
     @pytest.mark.flaky(reruns=3, reruns_delay=1)
     @pytest.mark.parametrize(

test/test_storage_map.py

@@ -350,6 +350,17 @@ def test_edges(self):
         edges_check = {(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)}
         assert edges == edges_check
 
+    def test_make_node(self):
+        td = TensorDict({"obs": torch.tensor([0])})
+        tree = Tree(node_data=td)
+        assert tree.node_data is not None
+
+        tree = Tree.make_node(data=td)
+        assert tree.node_data is not None
+
+        tree = Tree.make_node(td)
+        assert tree.node_data is not None
+
 
 class TestMCTSForest:
     def dummy_rollouts(self) -> Tuple[TensorDict, ...]:

torchrl/_utils.py

@@ -18,7 +18,6 @@
 import warnings
 from contextlib import nullcontext
 from copy import copy
-from distutils.util import strtobool
 from functools import wraps
 from importlib import import_module
 from typing import Any, Callable, cast, TypeVar
@@ -35,6 +34,21 @@
 except ImportError:
     from torch._dynamo import is_compiling
 
+
+def strtobool(val: Any) -> bool:
+    """Convert a string representation of truth to a boolean.
+
+    True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values are 'n', 'no', 'f', 'false', 'off', and '0'.
+    Raises ValueError if 'val' is anything else.
+    """
+    val = val.lower()
+    if val in ("y", "yes", "t", "true", "on", "1"):
+        return True
+    if val in ("n", "no", "f", "false", "off", "0"):
+        return False
+    raise ValueError(f"Invalid truth value {val!r}")
+
+
 LOGGING_LEVEL = os.environ.get("RL_LOGGING_LEVEL", "INFO")
 logger = logging.getLogger("torchrl")
 logger.setLevel(getattr(logging, LOGGING_LEVEL))

torchrl/data/llm/__init__.py

@@ -11,7 +11,14 @@
 )
 from .prompt import PromptData, PromptTensorDictTokenizer
 from .reward import PairwiseDataset, RewardData
-from .utils import AdaptiveKLController, ConstantKLController, RolloutFromModel, LLMData, LLMOutput, LLMInput
+from .utils import (
+    AdaptiveKLController,
+    ConstantKLController,
+    LLMData,
+    LLMInput,
+    LLMOutput,
+    RolloutFromModel,
+)
 
 __all__ = [
     "AdaptiveKLController",

torchrl/data/llm/utils.py

@@ -543,8 +543,10 @@ def step_scheduler(self):
             while len(self._kl_queue):
                 self._kl_queue.remove(self._kl_queue[0])
 
+
 LLMInpOut = TypeVar("LLMInpOut")
 
+
 class LLMInput(TensorClass["nocast"]):
     """Represents the input to a Large Language Model (LLM).
 
@@ -557,11 +559,13 @@ class LLMInput(TensorClass["nocast"]):
     .. seealso:: :class:`~torchrl.data.LLMOutput` and :class:`~torchrl.data.LLMData`.
 
     """
+
     tokens: torch.Tensor
     attention_mask: torch.Tensor | None = None
     token_list: list[int] | list[list[int]] | None = None
     text: str | list[str] | None = None
 
+
 class LLMOutput(TensorClass["nocast"]):
     """Represents the output from a Large Language Model (LLM).
 
@@ -581,6 +585,7 @@ class LLMOutput(TensorClass["nocast"]):
     .. seealso:: :class:`~torchrl.data.LLMInput` and :class:`~torchrl.data.LLMData`.
 
     """
+
     tokens: torch.Tensor
     tokens_response: torch.Tensor | None = None
     token_list: list[int] | list[list[int]] | None = None
@@ -594,6 +599,7 @@ def from_vllm_output(cls: type[LLMInpOut], vllm_output) -> LLMInpOut:
         # placeholder
         raise NotImplementedError
 
+
 class LLMData(TensorClass["nocast"]):
     """Represents the input or output of a Large Language Model (LLM).
 
@@ -619,6 +625,7 @@ class LLMData(TensorClass["nocast"]):
     .. seealso:: :class:`~torchrl.data.LLMInput` and :class:`~torchrl.data.LLMOutput`.
 
     """
+
     tokens: torch.Tensor
     tokens_response: torch.Tensor | None = None
     attention_mask: torch.Tensor | None = None

torchrl/data/map/tree.py

@@ -122,7 +122,7 @@ def make_node(
         return cls(
             count=torch.zeros(()),
             wins=torch.zeros(()),
-            node=data.exclude("action", "next"),
+            node_data=data.exclude("action", "next"),
             rollout=rollout,
             subtree=subtree,
             device=device,

torchrl/envs/batched_envs.py

@@ -379,6 +379,14 @@ def __init__(
 
     is_spec_locked = EnvBase.is_spec_locked
 
+    def select_and_clone(self, name, tensor, selected_keys=None):
+        if selected_keys is None:
+            selected_keys = self._selected_step_keys
+        if name in selected_keys:
+            if self.device is not None and tensor.device != self.device:
+                return tensor.to(self.device, non_blocking=self.non_blocking)
+            return tensor.clone()
+
     @property
     def non_blocking(self):
         nb = self._non_blocking
@@ -1072,12 +1080,10 @@ def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
         selected_output_keys = self._selected_reset_keys_filt
 
         # select + clone creates 2 tds, but we can create one only
-        def select_and_clone(name, tensor):
-            if name in selected_output_keys:
-                return tensor.clone()
-
         out = self.shared_tensordict_parent.named_apply(
-            select_and_clone,
+            lambda *args: self.select_and_clone(
+                *args, selected_keys=selected_output_keys
+            ),
             nested_keys=True,
             filter_empty=True,
         )
@@ -1150,14 +1156,14 @@ def _step(
             # will be modified in-place at further steps
             device = self.device
 
-            def select_and_clone(name, tensor):
-                if name in self._selected_step_keys:
-                    return tensor.clone()
+            selected_keys = self._selected_step_keys
 
             if partial_steps is not None:
                 next_td = TensorDict.lazy_stack([next_td[i] for i in workers_range])
             out = next_td.named_apply(
-                select_and_clone, nested_keys=True, filter_empty=True
+                lambda *args: self.select_and_clone(*args, selected_keys),
+                nested_keys=True,
+                filter_empty=True,
             )
             if out_tds is not None:
                 out.update(
@@ -2010,20 +2016,8 @@ def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
         next_td = shared_tensordict_parent.get("next")
         device = self.device
 
-        if next_td.device != device and device is not None:
-
-            def select_and_clone(name, tensor):
-                if name in self._selected_step_keys:
-                    return tensor.to(device, non_blocking=self.non_blocking)
-
-        else:
-
-            def select_and_clone(name, tensor):
-                if name in self._selected_step_keys:
-                    return tensor.clone()
-
         out = next_td.named_apply(
-            select_and_clone,
+            self.select_and_clone,
             nested_keys=True,
             filter_empty=True,
             device=device,
@@ -2203,20 +2197,10 @@ def tentative_update(val, other):
         selected_output_keys = self._selected_reset_keys_filt
         device = self.device
 
-        if self.shared_tensordict_parent.device != device and device is not None:
-
-            def select_and_clone(name, tensor):
-                if name in selected_output_keys:
-                    return tensor.to(device, non_blocking=self.non_blocking)
-
-        else:
-
-            def select_and_clone(name, tensor):
-                if name in selected_output_keys:
-                    return tensor.clone()
-
         out = self.shared_tensordict_parent.named_apply(
-            select_and_clone,
+            lambda *args: self.select_and_clone(
+                *args, selected_keys=selected_output_keys
+            ),
             nested_keys=True,
             filter_empty=True,
             device=device,

torchrl/objectives/value/advantages.py

@@ -1281,8 +1281,18 @@ def __init__(
             skip_existing=skip_existing,
             device=device,
         )
-        self.register_buffer("gamma", torch.tensor(gamma, device=self._device))
-        self.register_buffer("lmbda", torch.tensor(lmbda, device=self._device))
+        self.register_buffer(
+            "gamma",
+            gamma.to(self._device)
+            if isinstance(gamma, Tensor)
+            else torch.tensor(gamma, device=self._device),
+        )
+        self.register_buffer(
+            "lmbda",
+            lmbda.to(self._device)
+            if isinstance(lmbda, Tensor)
+            else torch.tensor(lmbda, device=self._device),
+        )
         self.average_gae = average_gae
         self.vectorized = vectorized
         self.time_dim = time_dim