torchrl 0.11.0__cp314-cp314-win_amd64.whl

torchrl/testing/gym_helpers.py

@@ -0,0 +1,156 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Versioned gym environment name helpers for TorchRL tests."""
+
+from __future__ import annotations
+
+import sys
+
+from torchrl._utils import implement_for
+from torchrl.envs.libs.gym import _has_gym, gym_backend
+
+__all__ = [
+    "BREAKOUT_VERSIONED",
+    "CARTPOLE_VERSIONED",
+    "CLIFFWALKING_VERSIONED",
+    "HALFCHEETAH_VERSIONED",
+    "PENDULUM_VERSIONED",
+    "PONG_VERSIONED",
+]
+
+PYTHON_3_9 = sys.version_info.major == 3 and sys.version_info.minor <= 9
+
+# Module-level variables that will be set by _set_gym_environments
+_CARTPOLE_VERSIONED = None
+_HALFCHEETAH_VERSIONED = None
+_PENDULUM_VERSIONED = None
+_PONG_VERSIONED = None
+_BREAKOUT_VERSIONED = None
+_CLIFFWALKING_VERSIONED = None
+
+
+def CARTPOLE_VERSIONED():
+    """Return the versioned CartPole environment name for the current gym backend."""
+    if gym_backend() is not None:
+        _set_gym_environments()
+        return _CARTPOLE_VERSIONED
+
+
+def HALFCHEETAH_VERSIONED():
+    """Return the versioned HalfCheetah environment name for the current gym backend."""
+    if gym_backend() is not None:
+        _set_gym_environments()
+        return _HALFCHEETAH_VERSIONED
+
+
+def PONG_VERSIONED():
+    """Return the versioned Pong environment name for the current gym backend."""
+    # Gymnasium says that the ale_py behavior changes from 1.0
+    # but with python 3.12 it is already the case with 0.29.1
+    try:
+        import ale_py  # noqa: F401
+    except ImportError:
+        pass
+
+    if gym_backend() is not None:
+        _set_gym_environments()
+        return _PONG_VERSIONED
+
+
+def CLIFFWALKING_VERSIONED():
+    """Return the versioned CliffWalking environment name for the current gym backend."""
+    if gym_backend() is not None:
+        _set_gym_environments()
+        return _CLIFFWALKING_VERSIONED
+
+
+def BREAKOUT_VERSIONED():
+    """Return the versioned Breakout environment name for the current gym backend."""
+    # Gymnasium says that the ale_py behavior changes from 1.0
+    # but with python 3.12 it is already the case with 0.29.1
+    try:
+        import ale_py  # noqa: F401
+    except ImportError:
+        pass
+
+    if gym_backend() is not None:
+        _set_gym_environments()
+        return _BREAKOUT_VERSIONED
+
+
+def PENDULUM_VERSIONED():
+    """Return the versioned Pendulum environment name for the current gym backend."""
+    if gym_backend() is not None:
+        _set_gym_environments()
+        return _PENDULUM_VERSIONED
+
+
+def _set_gym_environments():
+    global _CARTPOLE_VERSIONED, _HALFCHEETAH_VERSIONED, _PENDULUM_VERSIONED, _PONG_VERSIONED, _BREAKOUT_VERSIONED, _CLIFFWALKING_VERSIONED
+
+    _CARTPOLE_VERSIONED = None
+    _HALFCHEETAH_VERSIONED = None
+    _PENDULUM_VERSIONED = None
+    _PONG_VERSIONED = None
+    _BREAKOUT_VERSIONED = None
+    _CLIFFWALKING_VERSIONED = None
+
+
+@implement_for("gym", None, "0.21.0")
+def _set_gym_environments():  # noqa: F811
+    global _CARTPOLE_VERSIONED, _HALFCHEETAH_VERSIONED, _PENDULUM_VERSIONED, _PONG_VERSIONED, _BREAKOUT_VERSIONED, _CLIFFWALKING_VERSIONED
+
+    _CARTPOLE_VERSIONED = "CartPole-v0"
+    _HALFCHEETAH_VERSIONED = "HalfCheetah-v2"
+    _PENDULUM_VERSIONED = "Pendulum-v0"
+    _PONG_VERSIONED = "Pong-v4"
+    _BREAKOUT_VERSIONED = "Breakout-v4"
+    _CLIFFWALKING_VERSIONED = "CliffWalking-v0"
+
+
+@implement_for("gym", "0.21.0", None)
+def _set_gym_environments():  # noqa: F811
+    global _CARTPOLE_VERSIONED, _HALFCHEETAH_VERSIONED, _PENDULUM_VERSIONED, _PONG_VERSIONED, _BREAKOUT_VERSIONED, _CLIFFWALKING_VERSIONED
+
+    _CARTPOLE_VERSIONED = "CartPole-v1"
+    _HALFCHEETAH_VERSIONED = "HalfCheetah-v4"
+    _PENDULUM_VERSIONED = "Pendulum-v1"
+    _PONG_VERSIONED = "ALE/Pong-v5"
+    _BREAKOUT_VERSIONED = "ALE/Breakout-v5"
+    _CLIFFWALKING_VERSIONED = "CliffWalking-v0"
+
+
+@implement_for("gymnasium", None, "1.0.0")
+def _set_gym_environments():  # noqa: F811
+    global _CARTPOLE_VERSIONED, _HALFCHEETAH_VERSIONED, _PENDULUM_VERSIONED, _PONG_VERSIONED, _BREAKOUT_VERSIONED, _CLIFFWALKING_VERSIONED
+
+    _CARTPOLE_VERSIONED = "CartPole-v1"
+    _HALFCHEETAH_VERSIONED = "HalfCheetah-v4"
+    _PENDULUM_VERSIONED = "Pendulum-v1"
+    _PONG_VERSIONED = "ALE/Pong-v5"
+    _BREAKOUT_VERSIONED = "ALE/Breakout-v5"
+    _CLIFFWALKING_VERSIONED = "CliffWalking-v0"
+
+
+@implement_for("gymnasium", "1.0.0", "1.1.0")
+def _set_gym_environments():  # noqa: F811
+    raise ImportError
+
+
+@implement_for("gymnasium", "1.1.0")
+def _set_gym_environments():  # noqa: F811
+    global _CARTPOLE_VERSIONED, _HALFCHEETAH_VERSIONED, _PENDULUM_VERSIONED, _PONG_VERSIONED, _BREAKOUT_VERSIONED, _CLIFFWALKING_VERSIONED
+
+    _CARTPOLE_VERSIONED = "CartPole-v1"
+    _HALFCHEETAH_VERSIONED = "HalfCheetah-v5"
+    _PENDULUM_VERSIONED = "Pendulum-v1"
+    _PONG_VERSIONED = "ALE/Pong-v5"
+    _BREAKOUT_VERSIONED = "ALE/Breakout-v5"
+    _CLIFFWALKING_VERSIONED = "CliffWalking-v1" if not PYTHON_3_9 else "CliffWalking-v0"
+
+
+if _has_gym:
+    _set_gym_environments()

torchrl/testing/llm_mocks.py

@@ -0,0 +1,119 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""Shared test fixtures and mock infrastructure for LLM tests."""
+from __future__ import annotations
+
+import random
+import string
+
+import torch
+
+
+class MockTransformerConfig:
+    """Mock config to mimic transformers model config."""
+
+    def __init__(self, vocab_size: int, max_position_embeddings: int = 2048):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = vocab_size  # For simplicity
+
+
+class MockTransformerOutput:
+    """Mock output that mimics transformers model output with dict-like access."""
+
+    def __init__(self, logits):
+        self.logits = logits
+
+    def __getitem__(self, key):
+        """Allow dict-like access for compatibility."""
+        if key == "logits":
+            return self.logits
+        raise KeyError(f"Key {key} not found in model output")
+
+
+class MockTransformerModel(torch.nn.Module):
+    """Mock transformer model that mimics the structure of HuggingFace models."""
+
+    def __init__(self, vocab_size: int, device: torch.device | str | int = "cpu"):
+        super().__init__()
+        device = torch.device(device)
+        self.config = MockTransformerConfig(vocab_size)
+        # Simple embedding layer that maps tokens to logits
+        self.embedding = torch.nn.Embedding(vocab_size, vocab_size, device=device)
+        self.device = device
+
+    def forward(self, input_ids, attention_mask=None, **kwargs):
+        """Forward pass that returns logits in the expected format."""
+        # Get embeddings (which we'll use as logits for simplicity)
+        logits = self.embedding(input_ids.to(self.device) % self.config.vocab_size)
+        # Return output object similar to transformers models
+        return MockTransformerOutput(logits)
+
+    def get_tokenizer(self):
+        from transformers import AutoTokenizer
+
+        return AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
+
+
+class DummyStrDataLoader:
+    """A dummy dataloader that yields random strings for LLM testing."""
+
+    def __init__(self, batch_size=0):
+        if isinstance(batch_size, tuple):
+            batch_size = torch.Size(batch_size).numel()
+        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 {"text": self.generate_random_string()}
+        else:
+            return {
+                "query": [self.generate_random_string() for _ in range(self.batch_size)]
+            }
+
+
+class DummyTensorDataLoader:
+    """A dummy dataloader that yields random token tensors for LLM testing."""
+
+    def __init__(self, batch_size=0, max_length=10, padding=False):
+        if isinstance(batch_size, tuple):
+            batch_size = torch.Size(batch_size).numel()
+        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)
+        rt = torch.randint(1, 10000, (length,))
+        return rt
+
+    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()
+            tokens = 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]
+                tokens = torch.stack(tensors)
+            else:
+                tokens = tensors
+        return {"tokens": tokens, "attention_mask": tokens != 0}