[Feature] LLMEnv and DataLoadingPrimer

ghstack-source-id: 0785280
Pull Request resolved: #2818

Author: vmoens

docs/source/reference/envs.rst

@@ -440,6 +440,7 @@ TorchRL offers a series of custom built-in environments.
     ChessEnv
     PendulumEnv
     TicTacToeEnv
+    LLMEnv
     LLMHashingEnv
 
 
@@ -1033,6 +1034,7 @@ to be able to create this other composition:
     Compose
     ConditionalSkip
     Crop
+    DataLoadingPrimer
     DTypeCastTransform
     DeviceCastTransform
     DiscreteActionProjection

test/test_env.py

@@ -12,6 +12,7 @@
 import pickle
 import random
 import re
+import string
 from collections import defaultdict
 from functools import partial
 from sys import platform
@@ -43,9 +44,11 @@
     CatTensors,
     ChessEnv,
     ConditionalSkip,
+    DataLoadingPrimer,
     DoubleToFloat,
     EnvBase,
     EnvCreator,
+    LLMEnv,
     LLMHashingEnv,
     ParallelEnv,
     PendulumEnv,
@@ -66,6 +69,7 @@
     TransformedEnv,
     UnsqueezeTransform,
 )
+from torchrl.envs.transforms.rlhf import as_padded_tensor
 from torchrl.envs.utils import (
     _StepMDP,
     _terminated_or_truncated,
@@ -97,7 +101,7 @@
 
 try:
     this_dir = os.path.dirname(os.path.realpath(__file__))
-    with open(os.path.join(this_dir, "configs", "atari.yaml"), "r") as file:
+    with open(os.path.join(this_dir, "configs", "atari.yaml")) as file:
         atari_confs = yaml.load(file, Loader=yaml.FullLoader)
     _atari_found = True
 except FileNotFoundError:
@@ -4526,6 +4530,141 @@ def test_skipping_history_env_collector(self, device_env, collector_cls):
             count += 1
 
 
+class TestLLMEnv:
+    @pytest.fixture(scope="class", autouse=True)
+    def set_capture(self):
+        with set_capture_non_tensor_stack(False):
+            yield None
+        return
+
+    class DummyDataLoader:
+        def __init__(self, batch_size=0):
+            self.batch_size = batch_size
+
+        def generate_random_string(self, length=10):
+            """Generate a random string of a given length."""
+            return "".join(random.choice(string.ascii_lowercase) for _ in range(length))
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            if self.batch_size == 0:
+                return self.generate_random_string()
+            else:
+                return [self.generate_random_string() for _ in range(self.batch_size)]
+
+    class DummyTensorDataLoader:
+        def __init__(self, batch_size=0, max_length=10, padding=False):
+            self.batch_size = batch_size
+            self.max_length = max_length
+            self.padding = padding
+
+        def generate_random_tensor(self):
+            """Generate a tensor of random int64 values."""
+            length = random.randint(1, self.max_length)
+            return torch.tensor(
+                [random.randint(0, 100) for _ in range(length)], dtype=torch.int64
+            )
+
+        def pad_tensor(self, tensor):
+            """Pad a tensor to the maximum length."""
+            padding_length = self.max_length - len(tensor)
+            return torch.cat((torch.zeros(padding_length, dtype=torch.int64), tensor))
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            if self.batch_size == 0:
+                tensor = self.generate_random_tensor()
+                return self.pad_tensor(tensor) if self.padding else tensor
+            else:
+                tensors = [
+                    self.generate_random_tensor() for _ in range(self.batch_size)
+                ]
+                if self.padding:
+                    tensors = [self.pad_tensor(tensor) for tensor in tensors]
+                    return torch.stack(tensors)
+                else:
+                    return tensors
+
+    @pytest.mark.parametrize(
+        "str2str,stack_method",
+        [
+            [True, None],
+            [False, "as_padded_tensor"],
+            # TODO: a bit experimental, fails with check_env_specs
+            # [False, "as_nested_tensor"],
+            [False, None],
+        ],
+    )
+    @pytest.mark.parametrize("batched", [True, False])
+    @pytest.mark.parametrize("device", [None, "cpu"])
+    def test_llm_env(self, str2str, batched, stack_method, device):
+        env = LLMEnv(str2str=str2str, device=device)
+        if str2str:
+            primer = DataLoadingPrimer(
+                dataloader=self.DummyDataLoader(),
+                data_keys=["observation"],
+                example_data="a string!",
+            )
+        else:
+            if stack_method is None:
+                stack_method = as_padded_tensor
+            primer = DataLoadingPrimer(
+                dataloader=self.DummyTensorDataLoader(padding=True),
+                data_keys=["observation"],
+                data_specs=[Unbounded(shape=(-1,), dtype=torch.int64)],
+                stack_method=stack_method,
+            )
+        assert not env.batch_locked
+        env = env.append_transform(primer)
+        assert not env.batch_locked
+        if batched:
+            td = env.reset(TensorDict(batch_size=[3]))
+            env.check_env_specs(break_when_any_done="both", tensordict=td)
+        else:
+            env.check_env_specs(break_when_any_done="both")
+
+    @pytest.mark.parametrize(
+        "str2str,stack_method",
+        [
+            [True, None],
+            [False, "as_padded_tensor"],
+            # TODO: a bit experimental, fails with check_env_specs
+            # [False, "as_nested_tensor"],
+            [False, None],
+        ],
+    )
+    @pytest.mark.parametrize("batched", [True, False])
+    @pytest.mark.parametrize("device", [None, "cpu"])
+    def test_llm_from_dataloader(self, str2str, batched, stack_method, device):
+        if str2str:
+            kwargs = {
+                "dataloader": self.DummyDataLoader(),
+                "data_keys": ["observation"],
+                "example_data": "a string!",
+            }
+        else:
+            if stack_method is None:
+                stack_method = as_padded_tensor
+            kwargs = {
+                "dataloader": self.DummyTensorDataLoader(padding=True),
+                "data_keys": ["observation"],
+                "data_specs": [Unbounded(shape=(-1,), dtype=torch.int64)],
+                "stack_method": stack_method,
+            }
+        kwargs.update({"str2str": str2str, "device": device})
+        env = LLMEnv.from_dataloader(**kwargs)
+        assert not env.batch_locked
+        if batched:
+            td = env.reset(TensorDict(batch_size=[3]))
+            env.check_env_specs(break_when_any_done="both", tensordict=td)
+        else:
+            env.check_env_specs(break_when_any_done="both")
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/data/tensor_specs.py

@@ -22,7 +22,6 @@
     Dict,
     Generic,
     List,
-    Optional,
     overload,
     Sequence,
     Tuple,

torchrl/envs/__init__.py

@@ -5,7 +5,7 @@
 
 from .batched_envs import ParallelEnv, SerialEnv
 from .common import EnvBase, EnvMetaData, make_tensordict
-from .custom import ChessEnv, LLMHashingEnv, PendulumEnv, TicTacToeEnv
+from .custom import ChessEnv, LLMEnv, LLMHashingEnv, PendulumEnv, TicTacToeEnv
 from .env_creator import env_creator, EnvCreator, get_env_metadata
 from .gym_like import default_info_dict_reader, GymLikeEnv
 from .libs import (
@@ -58,6 +58,7 @@
     Compose,
     ConditionalSkip,
     Crop,
+    DataLoadingPrimer,
     DeviceCastTransform,
     DiscreteActionProjection,
     DoubleToFloat,

torchrl/envs/common.py

@@ -9,7 +9,7 @@
 import warnings
 from copy import deepcopy
 from functools import partial, wraps
-from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
+from typing import Any, Callable, Iterator
 
 import numpy as np
 import torch

torchrl/envs/custom/__init__.py

@@ -4,6 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 from .chess import ChessEnv
-from .llm import LLMHashingEnv
+from .llm import LLMEnv, LLMHashingEnv
 from .pendulum import PendulumEnv
 from .tictactoeenv import TicTacToeEnv