TorchRL

https://docs.pytorch.org/rl
Admin
TorchRL is an open-source Reinforcement Learning library for PyTorch that provides efficient,...
Tokens:7,383,919
Snippets:52,677
Trust Score:10
Update:4 months ago
Show doc for...
Tokens:
Raw
### DQNTrainer Setup Example

Source: https://docs.pytorch.org/rl/main/reference/generated/torchrl.trainers.algorithms.DQNTrainer.html

This snippet demonstrates the initial setup for the DQNTrainer, including the creation of a Collector, DQNLoss module, ReplayBuffer, HardUpdate mechanism, and an optimizer. Ensure all necessary components like 'env', 'policy' are defined before running.

```python
from torchrl.collectors import Collector
from torchrl.objectives import DQNLoss
from torchrl.data import ReplayBuffer, LazyTensorStorage
from torchrl.objectives.utils import HardUpdate
from torch import optim

# Set up collector, loss, and replay buffer
collector = Collector(env, policy, frames_per_batch=128)

```

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

### LLMEnv Initialization Example

Source: https://docs.pytorch.org/rl/main/_modules/torchrl/envs/llm/envs.html

Demonstrates basic initialization of the LLMEnv with default settings. This setup is suitable for use with vLLM backends.

```python
from torchrl.envs import LLMEnv

env = LLMEnv()

```

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

### PPO Trainer Initialization Example

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

Instantiate and train the PPO trainer using a configuration object. Ensure the environment setup is configurable.

```python
config = PPOTrainerConfig(...)  # Configure with required parameters
trainer = instantiate(config)
trainer.train()
```

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

### PPOLoss Initialization and Forward Pass Example

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

Provides a complete, runnable example of initializing PPOLoss with a probabilistic actor and a value operator, and then performing a forward pass with sample data. This snippet includes all necessary imports and setup.

```python
>>> import torch
>>> from torch import nn
>>> from torchrl.data.tensor_specs 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.ppo import PPOLoss
>>> from tensordict import TensorDict
>>> n_act, n_obs = 4, 3
>>> spec = Bounded(-torch.ones(n_act), torch.ones(n_act), (n_act,))
>>> base_layer = nn.Linear(n_obs, 5)
>>> net = nn.Sequential(base_layer, nn.Linear(5, 2 * n_act), NormalParamExtractor())
>>> module = SafeModule(net, in_keys=["observation"], out_keys=["loc", "scale"])
>>> actor = ProbabilisticActor(
...     module=module,
...     distribution_class=TanhNormal,
...     in_keys=["loc", "scale"],
...     spec=spec)
>>> module = nn.Sequential(base_layer, nn.Linear(5, 1))
>>> value = ValueOperator(
...     module=module,
...     in_keys=["observation"])
>>> loss = PPOLoss(actor, value)
>>> batch = [2, ]
>>> action = spec.rand(batch)
>>> data = TensorDict({"observation": torch.randn(*batch, n_obs),
...         "action": action,
...         "action_log_prob": torch.randn_like(action[..., 1]),
...         ("next", "done"): torch.zeros(*batch, 1, dtype=torch.bool),
...         ("next", "terminated"): torch.zeros(*batch, 1, dtype=torch.bool),
...         ("next", "reward"): torch.randn(*batch, 1),
...         ("next", "observation"): torch.randn(*batch, n_obs),
...     }, batch)
>>> loss(data)
TensorDict(
    fields={

```

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

### Quick Example: Creating and Running Environments

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

Demonstrates how to create a single Gym environment, add transforms like RewardSum, and set up parallel environments with StepCounter. It also shows how to run a rollout on the parallel environment.

```python
from torchrl.envs import GymEnv, ParallelEnv, TransformedEnv
from torchrl.envs.transforms import RewardSum, StepCounter

# Create a single environment
env = GymEnv("Pendulum-v1")

# Add transforms
env = TransformedEnv(env, RewardSum())

# Create parallel environments
def make_env():
    return TransformedEnv(
        GymEnv("Pendulum-v1"),
        StepCounter(max_steps=200)
    )

parallel_env = ParallelEnv(4, make_env)

# Run a rollout
rollout = parallel_env.rollout(100)
```

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

### OpenMLEnv Initialization and Reset Example

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

Demonstrates how to initialize the OpenMLEnv with a specific dataset and batch size, and then reset the environment to get the initial observation, reward, and done status. This is useful for starting an RL episode.

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

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

### Check Transformers Installation

Source: https://docs.pytorch.org/rl/stable/_modules/torchrl/modules/llm/policies/vllm_wrapper.html

Ensures the 'transformers' library is installed. Raises an ImportError if it's not found, guiding the user to install it.

```python
# 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.
from __future__ import annotations

import collections
import importlib.util
import threading
import warnings
from typing import Any, Literal, TYPE_CHECKING

import torch
from tensordict import (
    lazy_stack,
    LazyStackedTensorDict,
    MetaData,
    NonTensorStack,
    set_list_to_stack,
    TensorDict,
    TensorDictBase,
)
from tensordict.tensorclass import from_dataclass, TensorClass
from tensordict.utils import _zip_strict, NestedKey
from torch import distributions as D
from torch.nn.utils.rnn import pad_sequence

from torchrl.envs.utils import _classproperty
from torchrl.modules.llm.policies.common import (
    _batching,
    _extract_responses_from_full_histories,
    ChatHistory,
    LLMWrapperBase,
    LogProbs,
    Masks,
    Text,
    Tokens,
)
from torchrl.modules.utils.utils import _unpad_tensors


_HAS_VLLM = importlib.util.find_spec("vllm") is not None
_HAS_TRANSFORMERS = importlib.util.find_spec("transformers") is not None

if TYPE_CHECKING:
    from vllm.inputs import TokensPrompt  # type: ignore[import-not-found]
    from vllm.outputs import RequestOutput  # type: ignore[import-not-found]
    from vllm.sampling_params import SamplingParams  # type: ignore[import-not-found]
elif _HAS_VLLM:
    from vllm.outputs import RequestOutput
    from vllm.sampling_params import SamplingParams

    try:
        from vllm.inputs import TokensPrompt
    except ImportError:
        # Fallback for older vLLM versions
        TokensPrompt = None
else:
    SamplingParams = None  # Will error at usage if vLLM not available
    RequestOutput = None
    TokensPrompt = None


def _require_transformers() -> None:
    if not _HAS_TRANSFORMERS:
        raise ImportError(
            "transformers is required for vLLMWrapper. Please install it with `pip install transformers`."
        )


def _require_vllm():
    """Import vLLM lazily.

    We intentionally avoid importing vLLM at module import time because importing vLLM can
    load native extensions that may hard-crash the interpreter on some platforms.
    """
    if not _HAS_VLLM:
        raise ImportError(
            "vllm is required for vLLMWrapper. Please install it with `pip install vllm`."
        )
    import vllm as _vllm  # local import is intentional / required

    return _vllm

```

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

### Install Dependencies

Source: https://docs.pytorch.org/rl/0.12/_downloads/e2a5193e019585d9cfa26d3eebfd8be3/coding_ppo.ipynb

Installs necessary libraries for running the tutorial, including torchrl, gym[mujoco], and tqdm.

```bash
!pip3 install torchrl
!pip3 install gym[mujoco]
!pip3 install tqdm

```

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

### ActorValueOperator Example Setup

Source: https://docs.pytorch.org/rl/stable/_modules/torchrl/modules/tensordict_module/actors.html

Illustrates the setup for an ActorValueOperator, which combines a common observation embedding network with separate actor and critic (value) operators. This example shows the necessary imports and the conceptual structure.

```python
>>> import torch
>>> from tensordict import TensorDict
>>> from torchrl.modules import ProbabilisticActor, SafeModule
>>> from torchrl.modules import ValueOperator, TanhNormal, ActorValueOperator, NormalParamExtractor

```

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

### Example Trainer Setup

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

Demonstrates the creation and configuration of a trainer object using various components like environment creators, policies, loss modules, and replay buffers.

```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)
```

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

### PythonExecutorService Initialization Example

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/envs/llm/transforms/tools.html

Example of how to register and get an instance of the PythonExecutorService using the 'ray' backend.

```python
>>> # Register the service
>>> from torchrl.services import get_services
>>> services = get_services(backend="ray")
```

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

### RandomTruncationTransform Usage Example

Source: https://docs.pytorch.org/rl/stable/_modules/torchrl/envs/transforms/transforms.html

This example demonstrates how to use RandomTruncationTransform within a TransformedEnv. It shows the setup with StepCounter and RandomTruncationTransform, and how to perform a rollout.

```python
from torchrl.envs import GymEnv, TransformedEnv, StepCounter
base_env = GymEnv("Pendulum-v1")
env = TransformedEnv(
    base_env,
    Compose(
        StepCounter(),
        RandomTruncationTransform(
            prob=0.1, min_horizon=50, max_horizon=200
        ),
    ),
)
rollout = env.rollout(300)
# Episode length will be at most 200 steps
print(rollout.shape)
```

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

### Install Dependencies

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

Installs the necessary libraries for the tutorial: torchrl, vmas, and tqdm. Run this in a Google Colab environment.

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

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

### Install Dependencies

Source: https://docs.pytorch.org/rl/0.12/_downloads/f3b53654a2b96b82010dd96f73ff5151/multiagent_competitive_ddpg.ipynb

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

```bash
!pip3 install torchrl
!pip3 install vmas
!pip3 install pettingzoo[mpe]==1.24.3
!pip3 install tqdm
```

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

### MultiSyncCollector Setup Example

Source: https://docs.pytorch.org/rl/main/_modules/torchrl/collectors/_multi_sync.html

Example demonstrating how to initialize a MultiSyncCollector with environment makers and a policy. Ensure to use a main guard for multiprocessing.

```python
from torchrl.envs.libs.gym import GymEnv
from tensordict.nn import TensorDictModule
from torch import nn
from torchrl.collectors import MultiSyncCollector
if __name__ == "__main__":
    env_maker = lambda: GymEnv("Pendulum-v1", device="cpu")
    policy = TensorDictModule(nn.Linear(3, 1), in_keys=["observation"], out_keys=["action"])
    collector = MultiSyncCollector(
        create_env_fn=[env_maker, env_maker],
```

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

### Download Full Example Code

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

Download the complete code for this example to run it locally.

```bash
wget https://raw.githubusercontent.com/pytorch/examples/main/multiagent_competitive_ddpg/multiagent_competitive_ddpg.py
```

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

### MultiAsyncCollector Example

Source: https://docs.pytorch.org/rl/main/reference/generated/torchrl.collectors.MultiSyncCollector.html

Demonstrates how to set up and run a MultiAsyncCollector for data collection. It initializes environments, a policy, a shared replay buffer, and then starts the collector asynchronously, updating a progress bar based on the replay buffer's write count until a target number of frames is reached. Ensure to set the gym backend and have ale_py installed.

```python
>>> from torchrl.modules import RandomPolicy
>>> 
>>> import time
>>> from functools import partial
>>> 
>>> import tqdm
>>> 
>>> from torchrl.collectors import MultiAsyncCollector
>>> from torchrl.data import LazyTensorStorage, ReplayBuffer
>>> from torchrl.envs import GymEnv, set_gym_backend
>>> import ale_py
>>> 
>>> # Set the gym backend to gymnasium
>>> set_gym_backend("gymnasium").set()
>>> 
>>> if __name__ == "__main__":
...     # Create a random policy for the Pong environment
...     env_fn = partial(GymEnv, "ALE/Pong-v5")
...     policy = RandomPolicy(env_fn().action_spec)
... 
...     # Initialize a shared replay buffer
...     rb = ReplayBuffer(storage=LazyTensorStorage(10000), shared=True)
... 
...     # Create a multi-async data collector with 16 environments
...     num_envs = 16
...     collector = MultiAsyncCollector(
...         [env_fn] * num_envs,
...         policy=policy,
...         replay_buffer=rb,
...         frames_per_batch=num_envs * 16,
...         total_frames=-1,
...     )
... 
...     # Progress bar to track the number of collected frames
...     pbar = tqdm.tqdm(total=100_000)
... 
...     # Start the collector asynchronously
...     collector.start()
... 
...     # Track the write count of the replay buffer
...     prec_wc = 0
...     while True:
...         wc = rb.write_count
...         c = wc - prec_wc
...         prec_wc = wc
... 
...         # Update the progress bar
...         pbar.update(c)
...         pbar.set_description(f"Write Count: {rb.write_count}")
... 
...         # Check the write count every 0.5 seconds
...         time.sleep(0.5)
... 
...         # Stop when the desired number of frames is reached
...         if rb.write_count . 100_000:
...             break
... 
...     # Shut down the collector
...     collector.async_shutdown()
```

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

### Starting MultiAsyncCollector for Data Collection

Source: https://docs.pytorch.org/rl/stable/reference/generated/torchrl.collectors.MultiAsyncCollector.html

Illustrates the setup and start of a MultiAsyncCollector for asynchronous data collection. This includes initializing environments, a policy, a shared replay buffer, and the collector itself, followed by starting the collection process.

```python
>>> from torchrl.modules import RandomPolicy
>>> import time
>>> from functools import partial
>>> 
>>> import tqdm
>>> 
>>> from torchrl.collectors import MultiAsyncCollector
>>> from torchrl.data import LazyTensorStorage, ReplayBuffer
>>> from torchrl.envs import GymEnv, set_gym_backend
>>> import ale_py
>>> 
>>> # Set the gym backend to gymnasium
>>> set_gym_backend("gymnasium").set()
>>> 
>>> if __name__ == "__main__":
...     # Create a random policy for the Pong environment
...     env_fn = partial(GymEnv, "ALE/Pong-v5")
...     policy = RandomPolicy(env_fn().action_spec)
... 
...     # Initialize a shared replay buffer
...     rb = ReplayBuffer(storage=LazyTensorStorage(10000), shared=True)
... 
...     # Create a multi-async data collector with 16 environments
...     num_envs = 16
...     collector = MultiAsyncCollector(
...         [env_fn] * num_envs,
...         policy=policy,
...         replay_buffer=rb,
...         frames_per_batch=num_envs * 16,
...         total_frames=-1,
...     )
... 
...     # Progress bar to track the number of collected frames
...     pbar = tqdm.tqdm(total=100_000)
... 
...     # Start the collector asynchronously
...
```

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

### Quick Example with SGLang Backend

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

Illustrates setting up an LLM Collector using the SGLang backend. This example shows creating an SGLang engine and connecting to a server, then configuring the policy and collector.

```python
from torchrl.modules.llm import SGLangWrapper, AsyncSGLang
from torchrl.envs.llm import ChatEnv
from torchrl.collectors.llm import LLMCollector

# Create SGLang engine (connects to server or launches managed server)
engine = AsyncSGLang.from_pretrained("Qwen/Qwen2.5-7B", tp_size=2)
# Or connect to existing server:
# engine = AsyncSGLang.connect("http://localhost:30000")
policy = SGLangWrapper(engine, tokenizer=tokenizer, input_mode="history")

# Create environment
env = ChatEnv(tokenizer=tokenizer)

# Create collector
collector = LLMCollector(env, policy, dialog_turns_per_batch=256)

```

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

### DistributionalQValueActor Example

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/modules/tensordict_module/actors.html

Demonstrates the setup of a DistributionalQValueActor using an MLP and ensuring log-softmax is applied to action values. Requires specific imports and TensorDict setup.

```python
>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule, TensorDictSequential
>>> from torch import nn
>>> from torchrl.data import OneHot
>>> from torchrl.modules import DistributionalQValueActor, MLP
>>> td = TensorDict({'observation': torch.randn(5, 4)}, [5])
>>> nbins = 3
>>> module = MLP(out_features=(nbins, 4), depth=2)
>>> # let us make sure that the output is a log-softmax
>>> module = TensorDictSequential(
...     TensorDictModule(module, ["observation"], ["action_value"]),
...     TensorDictModule(lambda x: x.log_softmax(-2), ["action_value"], ["action_value"]),
... )
>>> action_spec = OneHot(4)
```

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

### KLComputation Transform Example

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/envs/llm/transforms/kl.html

Demonstrates how to create and apply the KLComputation transform to a TensorDict. This example shows the setup for calculating KL divergence and adding it to the reward.

```python
from tensordict import TensorDict
import torch

# Create sample log-probs
gen_log_probs = torch.randn(2, 10)  # 2 samples, 10 tokens each
ref_log_probs = torch.randn(2, 10)

# Create data with next tensordict
next_td = TensorDict(
    {
        ("gen_log_probs", "full"): gen_log_probs,
        ("ref_log_probs", "full"): ref_log_probs,
        "reward": torch.randn(2, 10, 1),
    },
    batch_size=(2,)
)
data = TensorDict(next=next_td, batch_size=(2,)) # Incorrect: TensorDict does not accept a dict with a single key and batch_size

# Create KLComputation transform
kl_transform = KLComputation(
    gen_log_probs_full_key=("gen_log_probs", "full"),
    ref_log_probs_full_key=("ref_log_probs", "full"),
    kl_key="kl_penalty",
    add_to_reward=True,
    coef=1.0,
)

# Apply transform
result = kl_transform(data)
kl = result["next"].get("kl_penalty")
print(f"KL shape: {kl.shape}")
```

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

### Basic Trainer Setup and Training

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

This snippet demonstrates the basic setup and training process using the Trainer class. It shows how to initialize a trainer with a collector, total frames, loss module, and optimizer, and how to register hooks for updating weights and logging scalar values before starting the training loop.

```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()
```

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

### Post Initialization Setup

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

Ensures that version-specific reset and output transform functions are correctly bound within the appropriate Gym backend context.

```python
def _post_init(self):
        # writes the functions that are gym-version specific to the instance
        # once and for all. This is aimed at avoiding the need of decorating code
        # with set_gym_backend + allowing for parallel execution (which would
        # be troublesome when both an old version of gym and recent gymnasium
        # are present within the same virtual env).
        #
        # These calls seemingly do nothing but they actually get rid of the @implement_for decorator.
        # We execute them within the set_gym_backend context manager to make sure we get
        # the right implementation.
        #
        # This method is executed by the metaclass of GymWrapper.
        with set_gym_backend(self.get_library_name(self._env)):
            self._reset_output_transform = self._reset_output_transform
            self._output_transform = self._output_transform
```

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

### ReplayBufferTrainer Initialization Example

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

Example of initializing a ReplayBufferTrainer with a replay buffer and specifying batch size. This setup is used to manage data collection and sampling for training.

```python
        >>> rb_trainer = ReplayBufferTrainer(replay_buffer=replay_buffer, batch_size=N)
        >>> trainer.register_op("batch_process", rb_trainer.extend)
        >>> trainer.register_op("process_optim_batch", rb_trainer.sample)
        >>> trainer.register_op("post_loss", rb_trainer.update_priority)
```

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

### Quick Example with SGLang Backend

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

Illustrates setting up an LLMCollector with the SGLang backend. This involves creating an SGLang engine (either launching a managed server or connecting to an existing one), wrapping it with SGLangWrapper, defining a ChatEnv, and then initializing the LLMCollector.

```python
from torchrl.modules.llm import SGLangWrapper, AsyncSGLang
from torchrl.envs.llm import ChatEnv
from torchrl.collectors.llm import LLMCollector

# Create SGLang engine (connects to server or launches managed server)
engine = AsyncSGLang.from_pretrained("Qwen/Qwen2.5-7B", tp_size=2)
# Or connect to existing server:
# engine = AsyncSGLang.connect("http://localhost:30000")
policy = SGLangWrapper(engine, tokenizer=tokenizer, input_mode="history")

# Create environment
env = ChatEnv(tokenizer=tokenizer)

# Create collector
collector = LLMCollector(env, policy, dialog_turns_per_batch=256)
```

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

### CountdownEnv Reset Example

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

Demonstrates resetting the CountdownEnv to get an initial state.

```python
envgym = gym.make("DMC-cheetah-v0", transform=CatTensors(in_keys=["position", "velocity"], out_key="observation"))
envgym.reset()
```

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

### FinancialRegimeEnv Step Example

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

Demonstrates how to reset the environment, generate random actions, and step through the environment using the step and step_mdp methods.

```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)

```

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

### 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 settings.

```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"
```

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

### AsyncVLLM Initialization and LoadBalancer Example

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/modules/llm/backends/vllm/vllm_async.html

Demonstrates how to initialize an AsyncVLLM service and create a LoadBalancer with different strategy configurations.

```python
service = AsyncVLLM.from_pretrained("Qwen/Qwen2.5-3B", num_replicas=3)

# Simple strategy
lb = LoadBalancer(service, "requests")
actor_idx = lb.select_actor()

# Strategy hierarchy: try prefix-aware first, fall back to requests, then round-robin
lb = LoadBalancer(service, ["prefix-aware", "requests", "round-robin"])
actor_idx = lb.select_actor(prompt="Hello world")  # Uses prefix routing
actor_idx = lb.select_actor()  # Falls back to requests (no prompt)

# Custom configuration
lb = LoadBalancer(
    service,
    ["prefix-aware", "kv-cache"],
    prefix_length=16,
    overload_threshold=2.0
)
```

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

### DiffusionActor Initialization and Usage Example

Source: https://docs.pytorch.org/rl/main/_modules/torchrl/modules/tensordict_module/actors.html

Demonstrates how to initialize and use the DiffusionActor with a default score network. The example shows setting up the actor and passing a TensorDict to get the denoised action.

```python
import torch
from tensordict import TensorDict
from torchrl.modules import DiffusionActor

# obs_dim not required when using the default network
actor = DiffusionActor(action_dim=2, num_steps=10)
td = TensorDict({"observation": torch.randn(4, 3)}, batch_size=[4])
td = actor(td)
print(td["action"].shape)
```

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

### Instantiate and Rollout Meltingpot Environment

Source: https://docs.pytorch.org/rl/main/_modules/torchrl/envs/libs/meltingpot.html

Demonstrates how to get a Meltingpot configuration, build the environment, wrap it with MeltingpotWrapper, and perform a rollout for a specified number of steps. Requires the meltingpot library to be installed.

```python
import substrate
from torchrl.envs.libs.meltingpot import MeltingpotWrapper

substrate_config = substrate.get_config("commons_harvest__open")
mp_env = substrate.build_from_config(
    substrate_config, roles=substrate_config.default_player_roles
)
env_torchrl = MeltingpotWrapper(env=mp_env)
print(env_torchrl.rollout(max_steps=5))
```

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

### Pendulum Environment Setup

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

Initializes the Pendulum environment. Ensure PyTorch and Gymnasium are installed.

```python
import gymnasium as gym

env = gym.make("Pendulum-v1", render_mode="rgb_array")
observation, info = env.reset()

for _ in range(1000):
    action = env.action_space.sample()  # agent policy and action selection
    observation, reward, terminated, truncated, info = env.step(action)

    if terminated or truncated:
        observation, info = env.reset()
env.close()
```

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

### VTrace Example Usage with Value Network

Source: https://docs.pytorch.org/rl/stable/_modules/torchrl/objectives/value/advantages.html

Illustrates a basic setup for VTrace where a value network is defined to estimate state values. The example shows the creation of a TensorDictModule for the value network.

```python
import torch.nn as nn
from tensordict.nn import TensorDictModule

# Example of defining a value network:
# value_net = TensorDictModule(nn.Linear(3, 1), in_keys=["obs"], out_keys=["state_value"])

```

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

### Transport Backend Interface - Setup Connection and Weights

Source: https://docs.pytorch.org/rl/main/reference/collectors_weightsync.html

Illustrates the setup method for establishing a connection and applying weights on the receiver side.

```python
transport.setup_connection_and_weights_on_receiver(
    worker_idx=0,
    weights=buffer,
    model=policy,
    strategy=strategy,
)
```

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

### ImaginedEnv Example with Dummy World Model

Source: https://docs.pytorch.org/rl/main/_modules/torchrl/envs/model_based/imagined.html

Demonstrates how to create and use an ImaginedEnv with a dummy world model. This example shows the setup for wrapping a base environment and collecting an imagined rollout.

```python
import torch
from tensordict import TensorDict
from tensordict.nn import TensorDictModule
from torchrl.envs.model_based import ImaginedEnv, ModelBasedEnvBase
from torchrl.data import Composite, Unbounded
from torchrl.envs import GymEnv
from torchrl.policies import RandomPolicy

# A toy world model that returns zero-mean, identity covariance
class DummyWorldModel(torch.nn.Module):
    def __init__(self, obs_dim):
        super().__init__()
        self.obs_dim = obs_dim
    def forward(self, action, observation):
        # Assuming observation comes in as a dict with a "mean" key
        mean = observation.get("mean", observation)
        var = torch.eye(self.obs_dim).expand(*mean.shape[:-1], -1, -1)
        return mean, var

base_env = GymEnv("Pendulum-v1")
obs_dim = base_env.observation_spec["observation"].shape[-1]

wm = TensorDictModule(
    DummyWorldModel(obs_dim),
    in_keys=["action", "observation"],
    out_keys=[("next", "observation", "mean"), ("next", "observation", "var")],
)

imagined_env = ImaginedEnv(wm, base_env, next_observation_key=("next", "observation"))

# Collect an imagined rollout
rollout = imagined_env.rollout(max_steps=5, policy=RandomPolicy(imagined_env.action_spec))
```

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

### Quick Example with vLLM Backend

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

Demonstrates setting up an LLMCollector using the vLLM backend for asynchronous inference. Requires initializing the vLLM engine, wrapping it with vLLMWrapper, creating a ChatEnv, and then initializing the LLMCollector.

```python
from torchrl.modules.llm import vLLMWrapper, AsyncVLLM
from torchrl.envs.llm import ChatEnv
from torchrl.collectors.llm import LLMCollector

# Create vLLM engine
engine = AsyncVLLM.from_pretrained("Qwen/Qwen2.5-7B", num_replicas=2)
policy = vLLMWrapper(engine, input_mode="history")

# Create environment
env = ChatEnv(tokenizer=tokenizer)

# Create collector
collector = LLMCollector(env, policy, dialog_turns_per_batch=256)
```

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

### RoboHiveEnv Usage Example

Source: https://docs.pytorch.org/rl/stable/_modules/torchrl/envs/libs/robohive.html

Demonstrates how to instantiate and use the RoboHiveEnv. Requires RoboHive and its dependencies to be installed.

```python
from torchrl.envs import RoboHiveEnv

# Access the first available RoboHive environment
env_name = RoboHiveEnv.available_envs[0]

# Instantiate the environment
env = RoboHiveEnv(env_name)

# Perform a rollout for 3 steps
rollout_result = env.rollout(3)
print(rollout_result)
```

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

### Connect to SGLang Server and Initialize Wrapper

Source: https://docs.pytorch.org/rl/stable/_modules/torchrl/modules/llm/policies/sglang_wrapper.html

Demonstrates how to connect to an existing SGLang server or launch a managed server and then initialize the SGLangWrapper for text generation.

```python
from torchrl.modules.llm.backends.sglang import AsyncSGLang
from torchrl.modules.llm.policies import SGLangWrapper

# Connect to existing server
backend = AsyncSGLang.connect("http://localhost:30000")
wrapper = SGLangWrapper(backend, input_mode="text", generate=True)

# Or launch managed server

```

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

### LLMEnv Get Submodule Example

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

Retrieves a specific submodule from the environment using its fully-qualified string name.

```python
>>> module = env.get_submodule("net_b.net_c.conv")
```

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

### LLMEnv Get Parameter Example

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

Retrieves a specific parameter from the environment using its fully-qualified string name.

```python
>>> env.get_parameter("target.parameter.name")
Parameter(...)
```

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

### REDQLoss Initialization Example

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

Demonstrates how to initialize the REDQLoss module with actor and Q-value networks. Requires setting up networks and action/observation specs.

```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.redq import REDQLoss
>>> from tensordict import TensorDict
>>> n_act, n_obs = 4, 3
>>> obs_spec = torchrl.data.TensorSpec((n_obs,), dtype=torch.float32)
>>> action_spec = torchrl.data.TensorSpec(
...     shape=(n_act,),
...     dtype=torch.float32,
...     lower=torch.zeros(n_act),
...     upper=torch.ones(n_act),
...     dtype=torch.float32,
... ).expand(1, n_act)
>>> actor_net = SafeModule(
...     nn.Sequential(
...         nn.Linear(n_obs, 64),
...         nn.ReLU(),
...         nn.Linear(64, 64),
...         nn.ReLU(),
...         nn.Linear(64, n_act),
...         NormalParamExtractor(),
...     ),
...     in_keys=['observation'],
...     out_keys=['loc', 'scale'])
>>> actor = ProbabilisticActor(
...     module=actor_net,
...     in_keys=['observation'],
...     spec=obs_spec,
...     distribution_class=TanhNormal,
...     distribution_kwargs={
...         'min': action_spec.lower,
...         'max': action_spec.upper
...     },
...     return_log_normal=True)
>>> qvalue_net = SafeModule(
...     nn.Sequential(
...         nn.Linear(n_obs + n_act, 64),
...         nn.ReLU(),
...         nn.Linear(64, 64),
...         nn.ReLU(),
...         nn.Linear(64, 1),
...     ),
...     in_keys=['observation', 'action'],
...     out_keys=['state_action_value'])
>>> qvalue_net = ValueOperator(module=qvalue_net, in_keys=['observation', 'action'])
>>> loss_module = REDQLoss(actor, qvalue_net, num_qvalue_nets=3, sub_sample_len=2)
>>> print(loss_module)
```

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

### LLMEnv Get Buffer Example

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

Retrieves a specific buffer from the environment using its fully-qualified string name.

```python
>>> env.get_buffer("target.buffer.name")
torch.Tensor(...)
```

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

### Create and Configure a Trainer

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

Demonstrates the setup of a trainer with various components like collector, loss module, replay buffer, and logger. It also shows how to register operational steps for the trainer.

```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)
```

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

### Batching Computations

Source: https://docs.pytorch.org/rl/main/tutorials/pendulum.html

This snippet demonstrates the basic setup for batching computations. Ensure PyTorch is installed and imported.

```python
device=None,
is_shared=False)
```

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

### LLMEnv Initialization with Specific Parameters

Source: https://docs.pytorch.org/rl/main/_modules/torchrl/envs/llm/envs.html

Illustrates initializing LLMEnv with parameters such as vocabulary size, device, and attention mask handling. It also shows how to assign rewards and done states automatically.

```python
from torchrl.envs import LLMEnv
import torch

env = LLMEnv(
    vocab_size=32000,
    device=torch.device("cuda"),
    has_attention=True,
    assign_reward=True,
    assign_done=True,
    eos_token_id=101,
)

```

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

### Starting MultiAsyncCollector for Data Collection

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

Initiates asynchronous data collection using MultiAsyncCollector. This example demonstrates setting up a random policy, a shared replay buffer, and starting the collector to gather data from multiple environments.

```python
>>> from torchrl.modules import RandomPolicy
>>> 
>>>             >>> import time
>>> from functools import partial
>>> 
>>> import tqdm
>>> 
>>> from torchrl.collectors import MultiAsyncCollector
>>> from torchrl.data import LazyTensorStorage, ReplayBuffer
>>> from torchrl.envs import GymEnv, set_gym_backend
>>> import ale_py
>>> 
>>> # Set the gym backend to gymnasium
>>> set_gym_backend("gymnasium").set()
>>> 
>>> if __name__ == "__main__":
...     # Create a random policy for the Pong environment
...     env_fn = partial(GymEnv, "ALE/Pong-v5")
...     policy = RandomPolicy(env_fn().action_spec)
... 
...     # Initialize a shared replay buffer
...     rb = ReplayBuffer(storage=LazyTensorStorage(10000), shared=True)
... 
...     # Create a multi-async data collector with 16 environments
...     num_envs = 16
...     collector = MultiAsyncCollector(
...         [env_fn] * num_envs,
...         policy=policy,
...         replay_buffer=rb,
...         frames_per_batch=num_envs * 16,
...         total_frames=-1,
...     )
... 
...     # Progress bar to track the number of collected frames
...     pbar = tqdm.tqdm(total=100_000)
... 
...     # Start the collector asynchronously
...     collector.start()
... 

```

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

### Build and Rollout MeltingPot Environment

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

Demonstrates how to build a MeltingPot environment using substrate configuration and perform a rollout for a specified number of steps. Requires the `meltingpot` library to be installed.

```python
import meltingpot
import torch
from torchrl.envs.libs.meltingpot import MeltingpotWrapper

# Get substrate configuration
substrate_config = substrate.get_config("commons_harvest__open")

# Build the MeltingPot environment
mp_env = substrate.build_from_config(
    substrate_config, roles=substrate_config.default_player_roles
)

# Wrap the environment with MeltingpotWrapper
env_torchrl = MeltingpotWrapper(env=mp_env)

# Perform a rollout for 5 steps and print the result
print(env_torchrl.rollout(max_steps=5))
```

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

### LLMEnv Initialization Example

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

Demonstrates the initialization of the LLMEnv class with custom keys and parameters. This setup is useful for integrating with specific model architectures or data formats.

```python
from torchrl.envs import LLMEnv

llm_env = LLMEnv(
    token_key="input_ids",
    str_key="prompt",
    attention_key="attn_mask",
    action_key="generated_tokens",
    reward_key="reward",
    from_text=False,
    vocab_size=50278,
    has_attention=True,
    assign_reward=True,
    assign_done=True,
    eos_token_id=102,
)
print(llm_env)
```
Tokens:
Raw