TorchRL

https://docs.pytorch.org/rl/0.12
Admin
TorchRL is an open-source Reinforcement Learning library for PyTorch providing efficient, modular,...
Tokens:1,256,037
Snippets:10,258
Trust Score:10
Update:2 months ago
Show doc for...
Tokens:
Raw
### DdpgCnnQNet Initialization and Usage Example

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.modules.DdpgCnnQNet.html

Instantiates and demonstrates the usage of the DdpgCnnQNet. It shows how to create the network and pass a dummy observation and action to get a Q-value output.

```python
>>> from torchrl.modules import DdpgCnnQNet
>>> import torch
>>> net = DdpgCnnQNet()
>>> print(net)
DdpgCnnQNet(
  (convnet): ConvNet(
    (0): LazyConv2d(0, 32, kernel_size=(8, 8), stride=(4, 4))
    (1): ELU(alpha=1.0)
    (2): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2))
    (3): ELU(alpha=1.0)
    (4): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (5): ELU(alpha=1.0)
    (6): AdaptiveAvgPool2d(output_size=(1, 1))
    (7): Squeeze2dLayer()
  )
  (mlp): MLP(
    (0): LazyLinear(in_features=0, out_features=200, bias=True)
    (1): ELU(alpha=1.0)
    (2): Linear(in_features=200, out_features=200, bias=True)
    (3): ELU(alpha=1.0)
    (4): Linear(in_features=200, out_features=1, bias=True)
  )
)
>>> obs = torch.zeros(1, 3, 64, 64)
>>> action = torch.zeros(1, 4)
>>> value = net(obs, action)
>>> print(value.shape)
torch.Size([1, 1])
```

--------------------------------

### GPWorldModel Initialization and Fit Example

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/modules/models/gp.html

Demonstrates how to initialize the GPWorldModel and fit it with a dataset. Ensure botorch and gpytorch are installed.

```python
import torch
from tensordict import TensorDict

model = GPWorldModel(obs_dim=4, action_dim=1)
dataset = TensorDict(
    {
        "observation": torch.randn(50, 4),
        "action": torch.randn(50, 1),
        ("next", "observation"): torch.randn(50, 4),
    },
    batch_size=[50],
)
model.fit(dataset)
```

--------------------------------

### Install Dependencies

Source: https://docs.pytorch.org/rl/0.12/_downloads/a977047786179278d12b52546e1c0da8/multiagent_ppo.ipynb

Install necessary libraries for the tutorial, including torchrl, vmas, and tqdm.

```bash
!pip3 install torchrl
!pip3 install vmas
!pip3 install tqdm

```

--------------------------------

### OpenMLEnv Initialization and Reset Example

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/envs/libs/openml.html

Demonstrates how to initialize an OpenMLEnv with a specific dataset and batch size, and then reset the environment to get the initial observation. Requires scikit-learn to be installed.

```python
env = OpenMLEnv("adult_onehot", batch_size=[2, 3])
print(env.reset())
```

--------------------------------

### Install Gymnasium with Atari Support

Source: https://docs.pytorch.org/rl/0.12/_sources/tutorials/torchrl_envs.rst.txt

Install the necessary packages for using Gymnasium environments with Atari support. This is a prerequisite for running the examples.

```bash
pip install "gymnasium[atari]" pygame
```

--------------------------------

### Install Dependencies

Source: https://docs.pytorch.org/rl/0.12/tutorials/export.html

Install necessary libraries for the tutorial, including tensordict, torchrl, and gymnasium with Atari support.

```bash
!pip install tensordict
!pip install torchrl
!pip install "gymnasium[atari]"
```

--------------------------------

### Instantiate GenesisEnv

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/envs/libs/genesis.html

Example of how to instantiate the GenesisEnv with a specific environment name. Ensure the 'genesis-world' package is installed.

```python
from torchrl.envs import GenesisEnv
env = GenesisEnv(env_name="franka_reach")
td = env.rollout(10)
```

--------------------------------

### AsyncVLLM Usage Examples

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.modules.llm.AsyncVLLM.html

Demonstrates how to initialize and use the AsyncVLLM service for text generation. Includes simple single-GPU setup, advanced multi-GPU tensor parallelism with multiple replicas, and direct initialization using AsyncEngineArgs.

```python
from torchrl.modules.llm import AsyncVLLM
from vllm import SamplingParams

# Simple usage - single GPU, single replica
service = AsyncVLLM.from_pretrained("Qwen/Qwen2.5-3B")

# Advanced usage - multi-GPU tensor parallel with multiple replicas
service = AsyncVLLM.from_pretrained(
    "Qwen/Qwen2.5-7B",
    num_devices=2,  # Use 2 GPUs for tensor parallelism
    num_replicas=2,  # Create 2 replicas for higher throughput
    max_model_len=4096
)

# Generate text
sampling_params = SamplingParams(temperature=0.7, max_tokens=100)
result = service.generate("Hello, world!", sampling_params)
print(result.outputs[0].text)

# Alternative: using AsyncEngineArgs directly for advanced configuration
from vllm import AsyncEngineArgs
engine_args = AsyncEngineArgs(
    model="Qwen/Qwen2.5-3B",
    tensor_parallel_size=2
)
service = AsyncVLLM.launch(engine_args, num_replicas=2)
```

--------------------------------

### Install Dependencies for TorchRL MARL Tutorial

Source: https://docs.pytorch.org/rl/0.12/_sources/tutorials/multiagent_ppo.rst.txt

Install necessary libraries for running the multi-agent PPO tutorial in Google Colab.

```bash
!pip3 install torchrl
!pip3 install vmas
!pip3 install tqdm
```

--------------------------------

### Install habitat-sim and habitat-lab

Source: https://docs.pytorch.org/rl/0.12/reference/generated/knowledge_base/HABITAT.html

Install habitat-sim with bullet physics in headless mode and then install the habitat-lab package from GitHub. Environment variables are set to reduce log verbosity.

```bash
conda install habitat-sim withbullet headless -c conda-forge -c aihabitat-nightly -y
pip install git+https://github.com/facebookresearch/habitat-lab.git#subdirectory=habitat-lab

# This is to reduce verbosity
export MAGNUM_LOG=quiet && export HABITAT_SIM_LOG=quiet
```

--------------------------------

### Basic Video Recording Setup

Source: https://docs.pytorch.org/rl/0.12/reference/generated/knowledge_base/VIDEO_CUSTOMISATION.html

Sets up a basic video recorder with a CSV logger and records a short rollout from a Gym environment. Ensure FFmpeg is installed and accessible.

```python
from torchrl.envs import GymEnv, TransformedEnv
from torchrl.record import CSVLogger, VideoRecorder

logger = CSVLogger(exp_name="my_exp")
env = GymEnv("CartPole-v1", from_pixels=True, pixels_only=False)

recorder = VideoRecorder(logger, tag="my_video")
record_env = TransformedEnv(env, recorder)
rollout = record_env.rollout(max_steps=3)
recorder.dump()

```

--------------------------------

### Quick Start: Collect Trajectory with GymEnv and QValueActor

Source: https://docs.pytorch.org/rl/0.12/_downloads/36fe09d5d4546649ee1a029c7144936e/torchrl_demo.ipynb

This snippet demonstrates how to quickly set up a reinforcement learning environment, define a policy, and collect a trajectory using TorchRL. It requires PyTorch and TorchRL to be installed.

```python
import torch
from torchrl.envs import GymEnv
from torchrl.modules import MLP, QValueActor

env = GymEnv("CartPole-v1")
actor = QValueActor(
    MLP(
        in_features=env.observation_spec["observation"].shape[-1],
        out_features=2,
        num_cells=[64, 64],
    ),
    in_keys=["observation"],
    spec=env.action_spec,
)
rollout = env.rollout(max_steps=200, policy=actor)
print(
    f"Collected {rollout.shape[0]} steps, total reward: {rollout['next', 'reward'].sum().item():.0f}"
)
```

--------------------------------

### Install TorchRL and Tensordict

Source: https://docs.pytorch.org/rl/0.12/tutorials/getting-started-2.html

Add these lines to the beginning of your notebook to install the necessary libraries.

```bash
!pip install tensordict
!pip install torchrl

```

--------------------------------

### Collector Quick Example

Source: https://docs.pytorch.org/rl/0.12/reference/collectors.html

A quick example demonstrating how to use the `Collector` class for data collection. This example sets up a batched environment and a collector, then iterates through the collector to gather data for training. It also shows how to update policy weights and shut down the collector.

```APIDOC
## Quick Example

```python
from torchrl.collectors import Collector
from torchrl.envs import GymEnv, ParallelEnv

# Create a batched environment
def make_env():
    return GymEnv("Pendulum-v1")

env = ParallelEnv(4, make_env)

# Create collector
collector = Collector(
    env,
    policy=my_policy,
    frames_per_batch=200,
    total_frames=10000,
)

# Collect data
for data in collector:
    # data is a TensorDict with shape [4, 50] (4 envs, 50 steps each)
    # Use data for training...

    # Update policy weights periodically
    if should_update:
        collector.update_policy_weights_()

collector.shutdown()
```
```

--------------------------------

### Initialize LLMEnv with DummyStrDataLoader

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.collectors.llm.LLMCollector.html

This example demonstrates how to set up an LLM environment using a dummy data loader, a tokenizer, and a vLLM policy wrapper. It configures the environment for text-based input and specifies batch size and grouping for repeated tokens.

```python
>>> import vllm
>>> from torchrl.modules import vLLMWrapper
>>> from torchrl.testing.mocking_classes import DummyStrDataLoader
>>> from torchrl.envs import LLMEnv
>>> llm_model = vllm.LLM("gpt2")
>>> tokenizer = llm_model.get_tokenizer()
>>> tokenizer.pad_token = tokenizer.eos_token
>>> policy = vLLMWrapper(llm_model)
>>> dataloader = DummyStrDataLoader(1)
>>> env = LLMEnv.from_dataloader(
...    dataloader=dataloader,
...    tokenizer=tokenizer,
...    from_text=True,
...    batch_size=1,
...    group_repeats=True,
... )
```

--------------------------------

### Get Default Multiprocessing Start Method

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/_utils.html

Returns TorchRL's preferred multiprocessing start method. Defaults to 'spawn' for safety and compatibility, but respects user-defined global start methods.

```python
def _get_default_mp_start_method() -> str:
    """Returns TorchRL's preferred multiprocessing start method.

    If the user has explicitly set a global start method via ``mp.set_start_method()``,
    that method is returned. Otherwise, defaults to ``"spawn"`` for improved safety
    across backends and to avoid known issues with ``fork`` in multi-threaded programs.
    """
    # Check if user has explicitly set a global start method
    try:
        current = mp.get_start_method(allow_none=True)
        if current is not None:
            return current
    except (TypeError, RuntimeError):
        pass
    return "spawn"
```

--------------------------------

### Environment Reset and Step Example

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.envs.EnvBase.html

Demonstrates resetting an environment and performing a series of steps, including action computation and environment updates using step_mdp.

```python
from torchrl.envs import GymEnv
env = GymEnv("Pendulum-1")
data = env.reset()
for i in range(10):
    # compute action
    env.rand_action(data)
    # Perform action
    next_data = env.step(reset_data)
    data = env.step_mdp(next_data)
```

--------------------------------

### Trainer Setup and Instantiation

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/trainers/helpers/trainers.html

Demonstrates the creation of a trainer instance using `make_trainer`, configuring various components like collector, loss module, replay buffer, and logger. It also shows how to print the trainer object.

```python
env_maker = EnvCreator(lambda: GymEnv("Pendulum-v0"))
env_proof = env_maker()
obs_spec = env_proof.observation_spec
action_spec = env_proof.action_spec
net = torch.nn.Linear(env_proof.observation_spec.shape[-1], action_spec.shape[-1])
net_value = torch.nn.Linear(env_proof.observation_spec.shape[-1], 1)  # for the purpose of testing
policy = SafeModule(action_spec, net, in_keys=["observation"], out_keys=["action"])
value = ValueOperator(net_value, in_keys=["observation"], out_keys=["state_action_value"])
collector = Collector(env_maker, policy, total_frames=100)
loss_module = DDPGLoss(policy, value, gamma=0.99)
recorder = env_proof
target_net_updater = None
policy_exploration = EGreedyWrapper(policy)
replay_buffer = TensorDictReplayBuffer()
dir = tempfile.gettempdir()
logger = TensorboardLogger(exp_name=dir)
trainer = make_trainer(collector, loss_module, recorder, target_net_updater, policy_exploration,
    replay_buffer, logger)
print(trainer)
```

--------------------------------

### LazyMemmapStorage Example: Setting and Getting Data

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.data.replay_buffers.LazyMemmapStorage.html

Illustrates how to set and get data from a LazyMemmapStorage. It shows how TensorDicts are stored and retrieved, with nested structures preserved.

```APIDOC
Examples
```
>>> data = TensorDict({
...     "some data": torch.randn(10, 11),
...     ("some", "nested", "data"): torch.randn(10, 11, 12),
... }, batch_size=[10, 11])
>>> storage = LazyMemmapStorage(100)
>>> storage.set(range(10), data)
>>> len(storage)  # only the first dimension is considered as indexable
10
>>> storage.get(0)
TensorDict(
    fields={
        some data: MemoryMappedTensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, is_shared=False),
        some: TensorDict(
            fields={
                nested: TensorDict(
                    fields={
                        data: MemoryMappedTensor(shape=torch.Size([11, 12]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([11]),
                    device=cpu,
                    is_shared=False)},
            batch_size=torch.Size([11]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([11]),
    device=cpu,
    is_shared=False)

```
```

--------------------------------

### Setup PPO Loss Module and Optimizer

Source: https://docs.pytorch.org/rl/0.12/_downloads/a977047786179278d12b52546e1c0da8/multiagent_ppo.ipynb

Initializes the PPO loss module using ClipPPOLoss, configures key mappings for data access, builds the GAE value estimator, and sets up an Adam optimizer.

```python
loss_module = ClipPPOLoss(
    actor_network=policy,
    critic_network=critic,
    clip_epsilon=clip_epsilon,
    entropy_coeff=entropy_eps,
    normalize_advantage=False,  # Important to avoid normalizing across the agent dimension
)
loss_module.set_keys(  # We have to tell the loss where to find the keys
    reward=env.reward_key,
    action=env.action_key,
    value=("agents", "state_value"),
    # These last 2 keys will be expanded to match the reward shape
    done=("agents", "done"),
    terminated=("agents", "terminated"),
)


loss_module.make_value_estimator(
    ValueEstimators.GAE, gamma=gamma, lmbda=lmbda
)  # We build GAE
GAE = loss_module.value_estimator

optim = torch.optim.Adam(loss_module.parameters(), lr)
```

--------------------------------

### worker_idx

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/collectors/_base.html

Get or set the worker index for a collector, useful in distributed setups.

```APIDOC
## worker_idx

### Description
Get or set the worker index for this collector. This is particularly useful when the collector is part of a distributed setup.

### Getter
- **Returns**: The worker index (0-indexed).
- **Raises**: `RuntimeError` if `worker_idx` has not been set.

### Setter
- **Args**:
  - `value` (int | None): The worker index (0-indexed) or None.
```

--------------------------------

### Basic Trainer Setup and Training

Source: https://docs.pytorch.org/rl/0.12/_sources/reference/trainers.rst.txt

Instantiate a Trainer with collector, loss module, and optimizer. Register hooks for custom training steps like weight updates and logging. Then, initiate the training process.

```python
from torchrl.trainers import Trainer
from torchrl.trainers import UpdateWeights, LogScalar

# Create trainer
trainer = Trainer(
    collector=collector,
    total_frames=1000000,
    loss_module=loss,
    optimizer=optimizer,
)

# Register hooks
UpdateWeights(collector, 10).register(trainer)
LogScalar("reward").register(trainer)

# Train
trainer.train()
```

--------------------------------

### Ray Backend Usage

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/collectors/_evaluator.html

Example of using the Evaluator with the Ray backend for distributed evaluation setups.

```APIDOC
## Ray backend

Typical usage -- **Ray backend**::

    evaluator = Evaluator(
        make_eval_env,
        policy_factory=make_eval_policy,
        max_steps=1000,
        backend="ray",
        init_fn=my_init_fn,
        num_gpus=1,
    )

    evaluator.trigger_eval(weights, step=step)
    result = evaluator.poll()
    evaluator.shutdown()
```

--------------------------------

### Basic Usage Example

Source: https://docs.pytorch.org/rl/0.12/reference/services.html

Demonstrates how to initialize the service registry, register a service, access it from another worker, and clean up.

```APIDOC
## Basic Usage Example

### Description
This example shows the fundamental steps of using the service registry: initializing the backend, getting the registry, registering a service, accessing it, and performing cleanup.

### Code
```python
import ray
from torchrl.services import get_services

# Initialize your backend (Ray in this example)
ray.init()

# Get the service registry for your chosen backend
services = get_services(backend="ray", namespace="my_namespace")

# Register a service (the class will become a distributed service)
services.register(
    "tokenizer",
    TokenizerService,
    vocab_size=50000,
    num_cpus=1,  # Backend-specific option (Ray)
)

# Access the service from any worker
# (other workers just need to call get_services with the same backend and namespace)
services = get_services(backend="ray", namespace="my_namespace")
tokenizer = services["tokenizer"]

# Call the service (syntax depends on backend)
# For Ray, you need to use .remote() and ray.get()
result = ray.get(tokenizer.encode.remote("Hello world"))

# Cleanup when done
services.reset()
ray.shutdown()
```
```

--------------------------------

### Get All Worker IDs

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/collectors/llm/weight_update/vllm.html

Returns a list containing only '0', indicating that this setup uses a single worker.

```python
return [0]
```

--------------------------------

### Get Available IsaacGym Environments

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/envs/libs/isaacgym.html

Retrieves a list of all available environment tasks supported by IsaacGym. This requires the `isaacgymenvs` library to be installed.

```python
import isaacgymenvs

available_envs = list(isaacgymenvs.tasks.isaacgym_task_map.keys())
```

--------------------------------

### Setup Video Recording Environment

Source: https://docs.pytorch.org/rl/0.12/_downloads/0edaef9fb710ac597eeeb399f4a3bd6e/getting-started-4.ipynb

Combine a Gym environment with VideoRecorder and a logger within a TransformedEnv to record videos.

```python
from torchrl.record import VideoRecorder

recorder = VideoRecorder(logger, tag="my_video")
record_env = TransformedEnv(env, recorder)
```

--------------------------------

### Import Dependencies for TorchRL

Source: https://docs.pytorch.org/rl/0.12/tutorials/multiagent_ppo.html

Import necessary libraries for running the TorchRL multi-agent PPO tutorial. This is a standard starting point for most TorchRL examples.

```python
import torch

from tensordict.nn import TensorDictModule
from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer
from torchrl.data.replay_buffer.samplers import SamplerWithoutReplacement
from torchrl.data.replay_buffer.storages import LazyMemmapStorage
from torchrl.envs import (Compose, DoubleToFloat,
                          ObservationNormalizer, StepCounter)
from torchrl.envs.libs.gym import GymEnv
from torchrl.modules import (Actor, ActorCritic, EGreedyWrapper, MLP, ProbabilisticActor, TanhNormal, ValueOperator)
from torchrl.objectives import PPO, ValueEstimators
from torchrl.objectives.common import LossModule
from torchrl.record import CSVLogger
from torchrl.trainers.helpers.utils import ExperimentConfig
```

--------------------------------

### Get Submodule Example

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.envs.llm.GSM8KRewardParser.html

Demonstrates how to retrieve a nested submodule using its fully-qualified string name. This method is efficient for checking submodule existence.

```python
A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

get_submodule("net_b.linear")
get_submodule("net_b.net_c.conv")
```

--------------------------------

### get_services with Custom Ray Initialization

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.services.get_services.html

Shows how to initialize the service registry with custom Ray configurations, including address and namespace.

```python
>>> # With custom Ray initialization
>>> services = get_services(
...     backend="ray",
...     ray_init_config={"address": "auto"},
...     namespace="my_experiment"
... )
```

--------------------------------

### Discovering Configuration Options with --help

Source: https://docs.pytorch.org/rl/0.12/reference/config.html

Use the --help flag with any TorchRL script to list all available configuration groups and their parameters. This is useful for exploring customization options.

```bash
python sota-implementations/ppo_trainer/train.py --help

```

--------------------------------

### Initialize RayEvalBackend

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/collectors/_evaluator.html

Initializes the Ray evaluation backend. Requires Ray to be installed and a policy factory to be provided. Handles setup for distributed evaluation workers.

```python
class _RayEvalBackend(_EvalBackend):
    """Wraps :class:`RayEvalWorker` to match the backend contract."""

    def __init__(
        self,
        env_maker: Callable[[], Any],
        policy_factory: Callable[..., Any] | None,
        max_steps: int | None,
        reward_keys: NestedKey,
        init_fn: Callable[[], None] | None,
        num_gpus: int,
        ray_kwargs: dict,
    ) -> None:
        if not _has_ray:
            raise RuntimeError(
                "Ray is required for backend='ray' but could not be found. "
                "Install it with: pip install ray"
            )
        if policy_factory is None:
            raise ValueError(
                "The 'ray' backend requires `policy_factory` (a callable "
                "`(env) -> policy`) because the policy is created inside "
                "the Ray actor process."
            )
        from torchrl.collectors.distributed import RayEvalWorker

        self._worker = RayEvalWorker(
            init_fn=init_fn,
            env_maker=env_maker,
            policy_maker=policy_factory,
            num_gpus=num_gpus,
            reward_keys=
            reward_keys
            if isinstance(reward_keys, tuple)
            else (reward_keys,),
            **ray_kwargs,
        )
        self._max_steps = max_steps
        self._last_step: int | None = None
        self._pending_flag = False
```

--------------------------------

### Setup Environment with Logging and Video Recording

Source: https://docs.pytorch.org/rl/0.12/_downloads/d04ba01cba3add54c2b1a3c96ae6320d/getting-started-5.ipynb

Configure the environment to log training data and record videos using CSVLogger and VideoRecorder. Ensure the environment is set up to capture pixel data if needed.

```python
from torchrl._utils import logger as torchrl_logger
from torchrl.record import CSVLogger, VideoRecorder

path = "./training_loop"
logger = CSVLogger(exp_name="dqn", log_dir=path, video_format="mp4")
video_recorder = VideoRecorder(logger, tag="video")
record_env = TransformedEnv(
    GymEnv("CartPole-v1", from_pixels=True, pixels_only=False), video_recorder
)
```

--------------------------------

### vLLM Weight Synchronization Flow Example (Ray Backend)

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/weight_update/llm/vllm_nccl.html

Illustrates the communication flow between a trainer and vLLM workers for weight synchronization using Ray for RPC and NCCL for collective data transfer. This example shows the signaling for collective start, the NCCL broadcast, and optional confirmation.

```text
    Trainer (rank 0)                    vLLM Workers (ranks 1+)
    ================                    =======================

    # 1. RPC: Signal collective start
    trainer.update_weights() ---------> [Ray RPC] --------> receiver.init_all_workers_group()
                                                            "I'm ready for collective"

    # 2. Collective: Broadcast weights
    NCCL broadcast -------------------- [GPU-GPU] ---------> NCCL receive
    (high bandwidth)                                        (parallel)

    # 3. RPC: Confirmation (optional)
    "broadcast done" <----------------- [Ray RPC] --------- "weights applied"
```

--------------------------------

### Initialize Ray and Get Service Registry

Source: https://docs.pytorch.org/rl/0.12/_sources/reference/services.rst.txt

Initializes the Ray backend and retrieves the service registry for a specific namespace. This is the starting point for using the service registry.

```python
import ray
from torchrl.services import get_services

# Initialize your backend (Ray in this example)
ray.init()

# Get the service registry for your chosen backend
services = get_services(backend="ray", namespace="my_namespace")

# Register a service (the class will become a distributed service)
services.register(
    "tokenizer",
    TokenizerService,
    vocab_size=50000,
    num_cpus=1,  # Backend-specific option (Ray)
)

# Access the service from any worker
# (other workers just need to call get_services with the same backend and namespace)
services = get_services(backend="ray", namespace="my_namespace")
tokenizer = services["tokenizer"]

# Call the service (syntax depends on backend)
# For Ray, you need to use .remote() and ray.get()
result = ray.get(tokenizer.encode.remote("Hello world"))

# Cleanup when done
services.reset()
ray.shutdown()
```

--------------------------------

### ToolRegistry Example

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.envs.llm.transforms.ToolRegistry.html

Demonstrates how to create a ToolRegistry, register a custom service, retrieve it by name, and execute it.

```python
>>> class AddService:
...     name = "add"
...     schema_in = {"a": int, "b": int}
...     schema_out = {"result": int}
...     def __call__(self, a, b, **kwargs):
...         return {"result": a + b}
>>> registry = ToolRegistry([AddService()])
>>> service = registry.get("add")
>>> result = service(a=1, b=2)
>>> print(result)
{"result": 3}
```

--------------------------------

### CQLLoss Initialization Example

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.objectives.CQLLoss.html

Demonstrates how to initialize the CQLLoss with actor and Q-value networks. Ensure that the actor and qvalue networks are properly defined and configured with their respective input and output keys.

```python
import torch
from torch import nn
from torchrl.data import Bounded
from torchrl.modules.distributions import NormalParamExtractor, TanhNormal
from torchrl.modules.tensordict_module.actors import ProbabilisticActor, ValueOperator
from torchrl.modules.tensordict_module.common import SafeModule
from torchrl.objectives.cql import CQLLoss
from tensordict import TensorDict

n_act, n_obs = 4, 3
spec = Bounded(-torch.ones(n_act), torch.ones(n_act), (n_act,))
net = nn.Sequential(nn.Linear(n_obs, 2 * n_act), NormalParamExtractor())
module = SafeModule(net, in_keys=["observation"], out_keys=["loc", "scale"])
actor = ProbabilisticActor(
    module=module,
    in_keys=["loc", "scale"],
    spec=spec,
    distribution_class=TanhNormal)
class ValueClass(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(n_obs + n_act, 1)
    def forward(self, obs, act):
        return self.linear(torch.cat([obs, act], -1))
module = ValueClass()
qvalue = ValueOperator(
    module=module,
    in_keys=['observation', 'action'])
loss = CQLLoss(actor, qvalue)
batch = [2, ]
action = spec.rand(batch)
data = TensorDict({
        "observation": torch.randn(*batch, n_obs),
        "action": action,
        ("next", "done"): torch.zeros(*batch, 1, dtype=torch.bool),
        ("next", "terminated"): torch.zeros(*batch, 1, dtype=torch.bool),
        ("next", "reward"): torch.randn(*batch, 1),

```

--------------------------------

### Initialize and Use OpenSpiel Wrapper

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/envs/libs/openspiel.html

Demonstrates how to initialize the OpenSpielWrapper with a game like 'chess', reset the environment, and take a step using a random action. The `return_state=True` argument is used to include the game state in the output.

```python
import pyspiel
from torchrl.envs import OpenSpielWrapper
from tensordict import TensorDict
base_env = pyspiel.load_game('chess').new_initial_state()
env = OpenSpielWrapper(base_env, return_state=True)
td = env.reset()
td = env.step(env.full_action_spec.rand())
print(td)
```

--------------------------------

### Setup Connection and Weights on Receiver

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/weight_update/_mp.html

Receives initial weights from the sender, applies them to the model, and starts a background receiver thread. Can be used with a specific worker index.

```python
def _setup_connection_and_weights_on_receiver_impl(
        self, *, worker_idx: int | None = None
    ) -> None:
        """Receive initial weights and start background receiver thread.

        This method:
        1. Receives initial weights from the sender via queue
        2. Applies them to the model
        3. Starts a background thread that listens for "receive" instructions

        Args:
            worker_idx: The worker index.
        """
        # Use stored worker_idx if not provided
        if worker_idx is None:
            worker_idx = self._worker_idx

        if worker_idx is None:
            raise RuntimeError(
                "worker_idx must be provided for _setup_connection_and_weights_on_receiver_impl."
            )

        # Receive initial weights from queue via transport
        if self._receiver_transport is None:
            raise RuntimeError("Receiver transport not set.")

        weights = self._receiver_transport.setup_connection_and_weights_on_receiver(
            worker_idx=worker_idx,
            weights=self.weights,
            model=self.model,
            strategy=self._strategy,
        )

        # Store received weights for later use
        if weights is not None:
            self._receiver_weights = weights

        # Apply weights to model
        if weights is not None and self.model is not None:
            self._strategy.apply_weights(self.model, weights, inplace=False)

        # Start background receiver thread
        self._start_background_receiver()
```

--------------------------------

### Initialize RayCollector

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/collectors/distributed/ray.html

Instantiate a RayCollector for distributed data collection. This example shows how to set up environment makers, a policy, and various collector parameters for a distributed setup.

```python
from torch import nn
from tensordict.nn import TensorDictModule
from torchrl.envs.libs.gym import GymEnv
from torchrl.collectors import Collector
from torchrl.collectors.distributed import RayCollector
env_maker = lambda: GymEnv("Pendulum-v1", device="cpu")
policy = TensorDictModule(nn.Linear(3, 1), in_keys=["observation"], out_keys=["action"])
distributed_collector = RayCollector(
    create_env_fn=[env_maker],
    policy=policy,
    collector_class=Collector,
    max_frames_per_traj=50,
    init_random_frames=-1,
    reset_at_each_iter=False,
    collector_kwargs={
        "device": "cpu",
        "storing_device": "cpu",
    },
    num_collectors=1,
    total_frames=10000,
    frames_per_batch=200,
)
for i, data in enumerate(collector):
    if i == 2:
        print(data)
        break
```

--------------------------------

### Show Available Configuration Options

Source: https://docs.pytorch.org/rl/0.12/_sources/reference/config.rst.txt

Use the --help flag with any TorchRL script to list all available configuration groups and their options.

```bash
python sota-implementations/ppo_trainer/train.py --help
```

--------------------------------

### PPOLoss Initialization and Usage Example

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/objectives/ppo.html

Demonstrates initializing PPOLoss and using it with a datacollector. It shows how the advantage can be computed by the loss function or externally.

```python
ppo_loss = PPOLoss(actor, critic)
advantage = GAE(critic)
data = next(datacollector)
losses = ppo_loss(data)
# equivalent
advantage(data)
losses = ppo_loss(data)
```

--------------------------------

### Get Available DMControl Environments

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/envs/libs/dm_control.html

Retrieves a list or dictionary of available environments from the dm_control suite. Raises an ImportError if dm_control is not installed. Useful for discovering and selecting environments.

```python
def _get_envs(to_dict: bool = True) -> dict[str, Any]:
    if not _has_dm_control:
        raise ImportError("Cannot find dm_control in virtual environment.")
    from dm_control import suite

    if not to_dict:
        return tuple(suite.BENCHMARKING) + tuple(suite.EXTRA)
    d = {}
    for tup in suite.BENCHMARKING:
        env_name = tup[0]
        d.setdefault(env_name, []).append(tup[1])
    for tup in suite.EXTRA:
        env_name = tup[0]
        d.setdefault(env_name, []).append(tup[1])
    return d.items()
```

--------------------------------

### Run Training Detached

Source: https://docs.pytorch.org/rl/0.12/reference/generated/knowledge_base/ISAACLAB.html

Starts the training process in detached mode, setting Weights & Biases logging to online. This command assumes the setup script has already been run.

```bash
steve step -d "$JOBID" 'WANDB_MODE=online bash /root/setup-and-run.sh'

```

--------------------------------

### Complete PPO Training Configuration Example

Source: https://docs.pytorch.org/rl/0.12/reference/config.html

A comprehensive configuration for PPO training, specifying environments, models, networks, collectors, optimizers, and loggers. This serves as a template for setting up training experiments.

```yaml
defaults:
  - env@training_env: batched_env
  - env@training_env.create_env_fn: gym
  - model@models.policy_model: tanh_normal
  - model@models.value_model: value
  - network@networks.policy_network: mlp
  - network@networks.value_network: mlp
  - collector: sync
  - replay_buffer: base
  - storage: tensor
  - sampler: without_replacement
  - writer: round_robin
  - trainer: ppo
  - optimizer: adam
  - loss: ppo
  - logger: wandb

# Network configurations
networks:
  policy_network:
    out_features: 2
    in_features: 4
    num_cells: [128, 128]

  value_network:
    out_features: 1
    in_features: 4
    num_calls: [128, 128]

# Model configurations
models:
  policy_model:
    network: ${networks.policy_network}
    in_keys: ["observation"]
    out_keys: ["action"]

  value_model:
    network: ${networks.value_network}
    in_keys: ["observation"]
    out_keys: ["state_value"]

# Environment
training_env:
  num_workers: 2
  create_env_fn:
    env_name: CartPole-v1
    _partial_: true

# Training components
trainer:
  collector: ${collector}
  optimizer: ${optimizer}
  loss_module: ${loss}
  logger: ${logger}
  total_frames: 100000

collector:
  create_env_fn: ${training_env}
  policy: ${models.policy_model}
  frames_per_batch: 1024

optimizer:
  lr: 0.001

loss:
  actor_network: ${models.policy_model}
  critic_network: ${models.value_model}

logger:
  exp_name: my_experiment

```

--------------------------------

### Running PPO Training with Command-Line Overrides

Source: https://docs.pytorch.org/rl/0.12/reference/config.html

Demonstrates how to run the PPO training script with default configurations or override specific parameters using command-line arguments. This allows for quick experimentation and parameter tuning.

```bash
# Use default configuration
python sota-implementations/ppo_trainer/train.py

# Override specific parameters
python sota-implementations/ppo_trainer/train.py optimizer.lr=0.0001

# Change environment
python sota-implementations/ppo_trainer/train.py training_env.create_env_fn.env_name=Pendulum-v1

# Use different collector
python sota-implementations/ppo_trainer/train.py collector=async

```

--------------------------------

### Copy and Run Setup Script

Source: https://docs.pytorch.org/rl/0.12/reference/generated/knowledge_base/ISAACLAB.html

Copies the setup script to the Steve job and executes it to build the environment. Ensure the script is named 'setup-and-run.sh'.

```bash
steve cp "$JOBID" ./setup-and-run.sh :/root/setup-and-run.sh
steve step "$JOBID" 'bash /root/setup-and-run.sh --build-only'

```

--------------------------------

### Starting a Multiprocess Collector Worker

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/collectors/_multi_base.html

This code snippet illustrates the setup and starting of a worker process for a multi-process collector. It includes extensive configuration options passed as keyword arguments to the `_main_async_collector` function. Pay attention to the error handling for common multiprocessing issues like non-serializable objects or tensors requiring gradients.

```python
kwargs = {
                    "policy_factory": policy_factory[i],
                    "pipe_child": pipe_child,
                    "queue_out": queue_out,
                    "create_env_fn": env_fun,
                    "create_env_kwargs": env_fun_kwargs,
                    "policy": policy_to_send,
                    "max_frames_per_traj": self.max_frames_per_traj,
                    "frames_per_batch": self.frames_per_batch_worker(worker_idx=i),
                    "reset_at_each_iter": self.reset_at_each_iter,
                    "policy_device": policy_device,
                    "storing_device": storing_device,
                    "env_device": env_device,
                    "exploration_type": self.exploration_type,
                    "reset_when_done": self.reset_when_done,
                    "idx": i,
                    "interruptor": self.interruptor,
                    "set_truncated": self.set_truncated,
                    "use_buffers": self._use_buffers,
                    "replay_buffer": self.replay_buffer,
                    "extend_buffer": self.extend_buffer,
                    "traj_pool": self._traj_pool,
                    "trust_policy": self.trust_policy,
                    "compile_policy": (
                        self.compiled_policy_kwargs if self.compiled_policy else False
                    ),
                    "cudagraph_policy": (
                        self.cudagraphed_policy_kwargs
                        if self.cudagraphed_policy
                        else False
                    ),
                    "no_cuda_sync": self.no_cuda_sync,
                    "collector_class": self.collector_class,
                    "postproc": (
                        self.postproc if self.replay_buffer is not None else None
                    ),
                    "weight_sync_schemes": self._weight_sync_schemes,
                    "worker_idx": i,  # Worker index for queue-based weight distribution
                    "init_random_frames": self.init_random_frames,
                    "profile_config": self._profile_config,
                    "trajs_per_batch": self._worker_trajs_per_batch,
                    "init_fn": self._worker_init_fn,
                }
                proc = _ProcessNoWarnCtx(
                    target=_main_async_collector,
                    num_threads=self.num_sub_threads,
                    _start_method=ctx.get_start_method(),
                    kwargs=kwargs,
                )
                # proc.daemon can't be set as daemonic processes may be launched by the process itself
                try:
                    proc.start()
                except TypeError as err:
                    if "cannot pickle" in str(err):
                        raise RuntimeError(
                            "A non-serializable object was passed to the collector workers."
                        ) from err
                except RuntimeError as err:
                    if "Cowardly refusing to serialize non-leaf tensor" in str(err):
                        raise RuntimeError(
                            "At least one of the tensors in the policy, replay buffer, environment constructor or postprocessor requires gradients. "
                            "This is not supported in multiprocessed data collectors.\n" 
                            "- For ReplayBuffer transforms, use a `transform_factory` instead with `delayed_init=True`.\n" 
                            "- Make sure your environment constructor does not reference tensors already instantiated on the main process.\n" 
                            "- Since no gradient can be propagated through the Collector pipes, the backward graph is never needed. Consider using detached tensors instead."
                        ) from err
                    elif "_share_fd_: only available on CPU" in str(
                        err
                    ) or "_share_filename_: only available on CPU" in str(err):
                        # This is a common failure mode on older PyTorch versions when using the
                        # "spawn" multiprocessing start method: the process object contains a

```

--------------------------------

### Get Full Done Spec

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.envs.llm.CountdownEnv.html

Retrieves the full specification for the 'done' signal, indicating termination conditions. This helps in understanding how the environment signals completion or truncation. The example uses a GymWrapper for Pendulum-v1.

```python
>>> import gymnasium
>>> from torchrl.envs import GymWrapper
>>> env = GymWrapper(gymnasium.make("Pendulum-v1"))
>>> env.full_done_spec
Composite(
    done: Categorical(
        shape=torch.Size([1]),
        space=DiscreteBox(n=2),
        device=cpu,
        dtype=torch.bool,
        domain=discrete),
    truncated:
```

--------------------------------

### SFTLoss Example with KL Regularization

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/objectives/llm/sft.html

This example demonstrates the end-to-end usage of SFTLoss for fine-tuning an LLM. It includes setting up chat data, tokenizers, models, training and reference policies, applying a transform to get reference log probabilities, and calculating the SFT loss with KL divergence regularization. Ensure all necessary imports are present and the data is prepared correctly.

```python
from torchrl.data.llm.chat import History, _CHAT_TEMPLATES
from torchrl.modules.llm import TransformersWrapper
from torchrl.objectives.llm.sft import SFTLoss
from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM
from tensordict import TensorDict, lazy_stack
import torch

# Create chat data
chats = [
    [
        {"role": "system", "content": "You are a helpful assistant.",},
        {"role": "user", "content": "Hello, how are you?",},
        {"role": "assistant", "content": "I'm doing well, thank you!",},
    ],
    [
        {"role": "system", "content": "You are a helpful assistant.",},
        {"role": "user", "content": "What's the weather like?",},
        {"role": "assistant", "content": "I can't check the weather for you.",},
    ],
]
history = History.from_chats(chats)

# Setup tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.chat_template = _CHAT_TEMPLATES["chatml_format"]
model = OPTForCausalLM(OPTConfig()).eval()

# Create training and reference policies
policy_train = TransformersWrapper(
    model,
    tokenizer=tokenizer,
    generate=False,
    from_text=True,
    chat_template_name="qwen",
)
policy_ref = TransformersWrapper(
    model,
    tokenizer=tokenizer,
    generate=False,
    from_text=True,
    return_log_probs=True,
    chat_template_name="qwen",
)

# Create the RetrieveLogProb transform
transform = RetrieveLogProb(
    policy_ref,
    assistant_only=True,
    tokenizer_kwargs={"chat_template_name": "qwen"},
    tokenizer=tokenizer,
)

# Prepare data
text = history[:, :-1].apply_chat_template(
    tokenizer=tokenizer, chat_template_name="qwen", add_generation_prompt=True
)
text_response = history.apply_chat_template(
    tokenizer=tokenizer, chat_template_name="qwen", add_generation_prompt=False
)
text_response = [
    txt[len(txt_start):] for txt, txt_start in zip(text_response, text)
]
td = TensorDict(
    text=text,
    text_response=text_response,
    history=history,
    next=TensorDict(
        reward=torch.randn(2, 1),
        done=torch.zeros(2, dtype=torch.bool),
        history=history,
    ),
    batch_size=(2,),
)
data = lazy_stack(list(td.unbind(0)))

# Apply the transform to get reference log probabilities
data = transform(data)
assert "ref_log_probs" in data["next"].keys()

# Use with SFTLoss for KL regularization
loss = SFTLoss(
    actor_network=policy_train,
    tokenizer=tokenizer,
    reduction="mean",
    normalize_by_seq_length=True,
    kl_to_ref_coeff=0.1,
    tokenizer_kwargs={"chat_template_name": "qwen"},
    loss_function="sft",
)
loss_vals = loss(data)
print(f"SFT Loss: {loss_vals.loss_sft.item():.4f}")
```

--------------------------------

### OpenSpielEnv Initialization and Usage

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/envs/libs/openspiel.html

Demonstrates how to initialize an OpenSpiel environment, interact with it using reset and step, and access environment properties. It also shows how to restore a specific environment state.

```APIDOC
## OpenSpielEnv

### Description
Initializes an OpenSpiel environment for a specified game.

### Parameters
- **game_string** (str) - Required - The name of the OpenSpiel game to load.
- **group_map** (MarlGroupMapType | dict[str, list[str]], optional) - Defines how agents are grouped in multi-agent settings. Defaults to MarlGroupMapType.ALL_IN_ONE_GROUP.
- **categorical_actions** (bool, optional) - If True, actions will be treated as categorical. Defaults to False.
- **return_state** (bool, optional) - If True, the environment will return the state in reset and step outputs. Defaults to False.

### Attributes
- **available_envs**: List of environments available to build.

### Examples
```python
from torchrl.envs import OpenSpielEnv
from tensordict import TensorDict

# Initialize environment
env = OpenSpielEnv("chess", return_state=True)
td = env.reset()

# Take a random step
td = env.step(env.full_action_spec.rand())
print(td)

# Access available environments
print(env.available_envs)

# Restore a specific state
td_restore = td["next"]
td = env.reset(td_restore)
print((td == td_restore).all())
```

### Representation
```python
print(f"{env.__class__.__name__}(env={env.game_string}, batch_size={env.batch_size}, device={env.device})")
```

## _build_env

### Description
Internal method to build the OpenSpiel environment.

### Parameters
- **game_string** (str) - The name of the OpenSpiel game.
- **requires_grad** (bool, optional) - Whether gradients are required. Defaults to False.
- **parameters** (optional) - Game-specific parameters.

### Raises
- **ImportError**: If open_spiel is not installed.
- **ValueError**: If unsupported kwargs are provided.
```

--------------------------------

### Get Multiprocessing Context

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/_utils.html

Returns a multiprocessing context with TorchRL's preferred start method. This function ensures consistent context allocation for primitives like Queues and Pipes across TorchRL components.

```python
def _get_mp_ctx(start_method: str | None = None):
    """Return a multiprocessing context with TorchRL's preferred start method.

    This is intentionally context-based (instead of relying on global
    ``mp.set_start_method``) so that TorchRL components can consistently allocate
    primitives (Queue/Pipe/Lock/Process) with a matching context.
    """
    if start_method is None:
        start_method = _get_default_mp_start_method()
    try:
        return mp.get_context(start_method)
    except ValueError:
        # Best effort fallback if a start method isn't supported on this platform.
        return mp.get_context("spawn")
```
Tokens:
Raw