### Full PPOLoss Example with PyTorch Modules

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

A complete example demonstrating the setup of actor and value networks using PyTorch modules and TensorDict, followed by the initialization and a forward pass of the PPOLoss module.

```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={

```

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

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

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

### RetrieveLogProb Transform Example

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

This example demonstrates how to set up a reference model and apply the RetrieveLogProb transform to obtain log-probabilities from chat history. It covers data preparation, model and tokenizer setup, transform instantiation, and applying the transform to get reference log probabilities.

```python
from torchrl.data.llm import History
from torchrl.modules.llm import TransformersWrapper
from torchrl.modules.llm.policies import ChatHistory
from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM
from tensordict import TensorDict, set_list_to_stack
import torch

# Set up list to stack for History
set_list_to_stack(True).set()

# 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)
print(f"Created history with shape: {history.shape}")

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

# Create reference model
ref_model = TransformersWrapper(
    model,
    tokenizer=tokenizer,
    input_mode="history",
    generate=False,
    return_log_probs=True,
    pad_output=True,
)

# Create the RetrieveLogProb transform
transform = RetrieveLogProb(
    ref_model,
    assistant_only=True,
    tokenizer=tokenizer,
)

# Prepare data using ChatHistory
chat_history = ChatHistory(full=history)
data = TensorDict(history=chat_history, batch_size=(2,)) # type: ignore

# Apply the transform to get reference log probabilities
result = transform(data)
log_probs_key = (ref_model.log_probs_key, "full")
ref_log_probs = result.get(log_probs_key)
print(f"Log-probs shape: {ref_log_probs.shape}")

```

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

### OpenXExperienceReplay Initialization

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.data.datasets.OpenXExperienceReplay.html

Demonstrates the initialization of OpenXExperienceReplay with specific configurations for samplers, writers, batch size, and transforms. This setup is typical for creating a replay buffer.

```python
OpenXExperienceReplay(
            is_shared=False),
            next: TensorDict(
                fields={
                    done: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.bool, is_shared=True),
                    info: TensorDict(
                        fields={
                            distance_from_origin: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            forward_reward: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                            qpos: MemoryMappedTensor(shape=torch.Size([1000000, 15]), device=cpu, dtype=torch.float64, is_shared=True),
                            qvel: MemoryMappedTensor(shape=torch.Size([1000000, 14]), device=cpu, dtype=torch.float64, is_shared=True),
                            reward_ctrl: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            reward_forward: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            reward_survive: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            success: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.bool, is_shared=True),
                            x_position: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            x_velocity: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            y_position: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            y_velocity: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True)},
                        batch_size=torch.Size([1000000]),
                        device=cpu,
                        is_shared=False),
                    observation: TensorDict(
                        fields={
                            achieved_goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                            desired_goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                            observation: MemoryMappedTensor(shape=torch.Size([1000000, 27]), device=cpu, dtype=torch.float64, is_shared=True)},
                        batch_size=torch.Size([1000000]),
                        device=cpu,
                        is_shared=False),
                    reward: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.float64, is_shared=True),
                    terminated: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.bool, is_shared=True),
                    truncated: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.bool, is_shared=True)},
                batch_size=torch.Size([1000000]),
                device=cpu,
                is_shared=False),
            observation: TensorDict(
                fields={
                    achieved_goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                    desired_goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                    observation: MemoryMappedTensor(shape=torch.Size([1000000, 27]), device=cpu, dtype=torch.float64, is_shared=True)},
                batch_size=torch.Size([1000000]),
                device=cpu,
                is_shared=False)},
        batch_size=32,
        transform=Compose(
        ),
        collate_fn=<function _collate_id at 0x120e21dc0>)
```

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

### Check for vLLM Installation

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

Checks if vLLM is installed. Raises an ImportError if not found, guiding the user to install it.

```python
import importlib.util

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

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

### 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,
... )
```

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

### Collector Start Method Example

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

Demonstrates how to start a collector asynchronously in a separate thread for decoupled data collection. Ensure a replay buffer is initialized before starting.

```python
from torchrl.modules import RandomPolicy
            import time
            from functools import partial
            import tqdm
            from torchrl.collectors import Collector
            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 = GymEnv("ALE/Pong-v5")
                policy = RandomPolicy(env.action_spec)
            
                # Initialize a shared replay buffer
                rb = ReplayBuffer(storage=LazyTensorStorage(1000), shared=True)
            
                # Create a synchronous data collector
                collector = Collector(
                    env,
                    policy=policy,
                    replay_buffer=rb,
                    frames_per_batch=256,
                    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()
```

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

### OpenXExperienceReplay Initialization and Usage

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.data.datasets.OpenXExperienceReplay.html

Demonstrates how to initialize OpenXExperienceReplay and iterate over its data, including reading from a stream and sampling with a dynamically defined batch size.

```APIDOC
>>> # Read data from a stream. Deliver entire trajectories when iterating
>>> dataset = OpenXExperienceReplay("cmu_stretch",
...     num_slices=num_slices, download=False, streaming=True)
>>> for data in dataset: # data does not have a consistent shape
...     break
>>> # Define batch-size dynamically
>>> data = dataset.sample(128)  # delivers 2 sub-trajectories of length 64
```

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

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

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

### Initialize OpenXExperienceReplay Dataset

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.data.datasets.OpenXExperienceReplay.html

Demonstrates how to initialize the OpenXExperienceReplay dataset, specifying parameters like batch size, number of slices, download status, streaming mode, and root directory. It then iterates through the dataset to print a batch, reshaping it according to the number of slices.

```python
from torchrl.data.datasets import OpenXExperienceReplay
import tempfile
# Download the data, and sample 128 elements in each batch out of two trajectories
num_slices = 2
with tempfile.TemporaryDirectory() as root:
    dataset = OpenXExperienceReplay("cmu_stretch", batch_size=128,
        num_slices=num_slices, download=True, streaming=False,
        root=root,
        )
    for batch in dataset:
        print(batch.reshape(num_slices, -1))
        break
```

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

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

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

### IQLLoss Example

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

Demonstrates the initialization and usage of the IQLLoss module with sample actor, Q-value, and value networks. It shows how to create a batch of data and compute the loss.

```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.iql import IQLLoss
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 QValueClass(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))
qvalue = SafeModule(
    QValueClass(),
    in_keys=["observation", "action"],
    out_keys=["state_action_value"],
)
value = SafeModule(
    nn.Linear(n_obs, 1),
    in_keys=["observation"],
    out_keys=["state_value"],
)
loss = IQLLoss(actor, qvalue, value)
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),
        ("next", "observation"): torch.randn(*batch, n_obs),
    }, batch)
loss(data)
```

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

### TD1Estimator Usage Example

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

This example demonstrates how to use the TD1Estimator module to get value estimates. It requires the `terminated` tensor to be defined.

```python
terminated = torch.zeros(1, 10, 1, dtype=torch.bool)
value_target = module(obs=obs, next_reward=reward, next_done=done, next_obs=next_obs, next_terminated=terminated)
```

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

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

```
```

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

### AsyncVLLM Initialization and Usage Example

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

Demonstrates how to initialize and use the AsyncVLLM service for text generation. Includes examples for simple and advanced configurations, as well as direct usage of 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)
```

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

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

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

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

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

### RandomTruncationTransform Example

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

This example demonstrates how to use RandomTruncationTransform to randomly truncate episodes. It shows the setup with GymEnv, StepCounter, and Compose, and then performs a rollout, printing the resulting shape.

```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)
torch.Size([...])
```

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

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

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

### Initialize and Sample from AtariDQNExperienceReplay with ReplayBufferEnsemble

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.data.datasets.AtariDQNExperienceReplay.html

Demonstrates how to initialize AtariDQNExperienceReplay datasets for different games and combine them using ReplayBufferEnsemble. It also shows how to sample from the ensemble buffer.

```python
from torchrl.data.datasets import AtariDQNExperienceReplay
from torchrl.data.replay_buffers import ReplayBufferEnsemble
# we change this parameter for quick experimentation, in practice it should be left untouched
AtariDQNExperienceReplay._max_runs = 2
dataset_asterix = AtariDQNExperienceReplay("Asterix/5", batch_size=128, slice_len=64, num_procs=4)
dataset_pong = AtariDQNExperienceReplay("Pong/5", batch_size=128, slice_len=64, num_procs=4)
dataset = ReplayBufferEnsemble(dataset_pong, dataset_asterix, batch_size=128, sample_from_all=True)
sample = dataset.sample()
print("first sample, Asterix", sample[0])
print("second sample, Pong", sample[1])
```

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

### InferenceServer and ThreadingTransport Setup

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

Demonstrates the basic setup of an InferenceServer using ThreadingTransport for in-process actors. This includes initializing the server with a policy, starting it, and obtaining a client handle for actors to send data.

```APIDOC
## InferenceServer with ThreadingTransport

### Description
This example shows how to set up an auto-batching inference server for use with actors running in the same process (e.g., as threads). It initializes the server with a policy and a `ThreadingTransport`, starts the server, and then obtains a client to send data to the server.

### Usage
```python
from tensordict.nn import TensorDictModule
from torchrl.modules.inference_server import (
    InferenceServer,
    ThreadingTransport,
)
import torch.nn as nn
import concurrent.futures

# Define the policy model
policy = TensorDictModule(
    nn.Sequential(nn.Linear(8, 64), nn.ReLU(), nn.Linear(64, 4)),
    in_keys=["observation"],
    out_keys=["action"],
)

# Initialize the transport and the inference server
transport = ThreadingTransport()
server = InferenceServer(policy, transport, max_batch_size=32)
server.start()

# Get the client handle from the transport
client = transport.client()

# Actors can now call client(td) to send data for inference.
# The server will automatically batch these requests.
with concurrent.futures.ThreadPoolExecutor(16) as pool:
    # Example of actor threads submitting tasks
    # future = pool.submit(client, sample_tensordict)
    pass

# Shut down the server when done
server.shutdown()
```

### Core Components
- **`InferenceServer(model, transport, max_batch_size)`**: Initializes the inference server with the model to serve and the transport mechanism. `max_batch_size` controls the maximum number of requests to batch.
- **`ThreadingTransport()`**: An in-process transport backend suitable for multi-threaded applications.
- **`transport.client()`**: Returns a client object that actors use to send data to the server.
```

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

### Instantiate and Sample OpenX Dataset

Source: https://docs.pytorch.org/rl/0.12/_modules/torchrl/data/datasets/openx.html

Demonstrates how to download the OpenX dataset, instantiate it with specific sampling parameters, and iterate over batches. Ensure the dataset is downloaded and configured for sampling sub-trajectories.

```python
from torchrl.data.datasets import OpenXExperienceReplay
import tempfile

# Download the data, and sample 128 elements in each batch out of two trajectories
num_slices = 2
with tempfile.TemporaryDirectory() as root:
    dataset = OpenXExperienceReplay("cmu_stretch", batch_size=128,
        num_slices=num_slices, download=True, streaming=False,
        root=root,
        )
    for batch in dataset:
        print(batch.reshape(num_slices, -1))
```

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

### Setup Inference Server with ThreadingTransport

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

Demonstrates the basic setup of an InferenceServer using ThreadingTransport for actors within the same process. Ensure the server is started before clients attempt to connect and shut it down when inference is complete.

```python
from tensordict.nn import TensorDictModule
from torchrl.modules.inference_server import (
    InferenceServer,
    ThreadingTransport,
)
import torch.nn as nn
import concurrent.futures

policy = TensorDictModule(
    nn.Sequential(nn.Linear(8, 64), nn.ReLU(), nn.Linear(64, 4)),
    in_keys=["observation"],
    out_keys=["action"],
)

transport = ThreadingTransport()
server = InferenceServer(policy, transport, max_batch_size=32)
server.start()
client = transport.client()

# actor threads call client(td) -- batched automatically
with concurrent.futures.ThreadPoolExecutor(16) as pool:
    ...

server.shutdown()
```

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

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

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

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

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

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

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

### 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]
```

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

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

```

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

### Get state_dict keys

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

This example shows how to retrieve the keys of the state dictionary for a module. It's useful for understanding what parameters and buffers are saved. The SKIP flag indicates that this example requires undefined variables to run.

```python
# xdoctest: +SKIP("undefined vars")
module.state_dict().keys()
```

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

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

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

### Setup Tokenizer and Model for KL Divergence

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

This example shows the necessary setup for a tokenizer and a causal language model, which are prerequisites for calculating KL divergence. Ensure the tokenizer's pad token is set correctly.

```python
from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM

# Setup tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
tokenizer.pad_token = tokenizer.eos_token
model = OPTForCausalLM(OPTConfig()).eval()
```

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

### Instantiate OpenMLExperienceReplay

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.data.datasets.OpenMLExperienceReplay.html

Demonstrates how to instantiate the OpenMLExperienceReplay dataset with specified samplers, writers, batch size, and transforms. This setup is typical for creating a replay buffer for reinforcement learning tasks.

```python
from torchrl.data.datasets import OpenMLExperienceReplay
from torchrl.data.replay_buffer import RandomSampler, ImmutableDatasetWriter
from torchrl.envs import Compose

# Assuming OpenMLExperienceReplay is initialized with some data source
# For demonstration, we'll show the instantiation parameters:
# dataset = OpenMLExperienceReplay(
#     samplers=RandomSampler,
#     writers=ImmutableDatasetWriter(),
#     batch_size=32,
#     transform=Compose(
#     ),
#     collate_fn=<function _collate_id at 0x120e21dc0>
# )

```

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

### Initialize and use StepCounter

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

Demonstrates how to wrap a base environment with StepCounter and perform a rollout. The output shows the 'step_count' and 'truncated' keys being populated.

```python
import gymnasium
from torchrl.envs import GymWrapper, TransformedEnv
from torchrl.envs.transforms import StepCounter

base_env = GymWrapper(gymnasium.make("Pendulum-v1"))
env = TransformedEnv(base_env,
    StepCounter(max_steps=5))
rollout = env.rollout(100)
print(rollout)
print(rollout["next", "step_count"])
```

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

### Multi-Step Actor Wrapper Example

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

Demonstrates wrapping an actor with MultiStepActorWrapper to handle multi-action sequences, often used with environments that require a lookahead window or macro-execution. This example shows setup with transformed environments and serial environments.

```python
import torch.nn
from torchrl.modules.tensordict_module.actors import MultiStepActorWrapper, Actor
from torchrl.envs import CatFrames, GymEnv, TransformedEnv, SerialEnv, InitTracker, Compose
from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod

time_steps = 6
n_obs = 4
n_action = 2
batch = 5

# Transforms a CatFrames in a stack of frames
def reshape_cat(data: torch.Tensor):
    return data.unflatten(-1, (time_steps, n_obs))
# an actor that reads `time_steps` frames and outputs one action per frame
# (actions are conditioned on the observation of `time_steps` in the past)
actor_base = Seq(
    Mod(reshape_cat, in_keys=["obs_cat"], out_keys=["obs_cat_reshape"]),
    Mod(torch.nn.Linear(n_obs, n_action), in_keys=["obs_cat_reshape"], out_keys=["action"])
)
# Wrap the actor to dispatch the actions
actor = MultiStepActorWrapper(actor_base, n_steps=time_steps)

env = TransformedEnv(
    SerialEnv(batch, lambda: GymEnv("CartPole-v1")),
    Compose(
        InitTracker(),
        CatFrames(N=time_steps, in_keys=["observation"], out_keys=["obs_cat"], dim=-1)
    )
)

print(env.rollout(100, policy=actor, break_when_any_done=False))
```

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

### 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),

```

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

### Get Data Collector

Source: https://docs.pytorch.org/rl/0.12/_downloads/25730743bad2ad4374b1a37c2e8d077a/coding_dqn.ipynb

Initializes and returns a data collector, choosing between Collector and MultiAsyncCollector based on the multiprocessing start method.

```python
def get_collector(
    stats,
    num_collectors,
    actor_explore,
    frames_per_batch,
    total_frames,
    device,
):
    # We can't use nested child processes with mp_start_method="fork"
    if is_fork:
        cls = Collector
        env_arg = make_env(parallel=True, obs_norm_sd=stats, num_workers=num_workers)
    else:
        cls = MultiAsyncCollector
        env_arg = [
            make_env(parallel=True, obs_norm_sd=stats, num_workers=num_workers)
        ] * num_collectors
    data_collector = cls(
        env_arg,
        policy=actor_explore,
        frames_per_batch=frames_per_batch,
        total_frames=total_frames,
        # this is the default behavior: the collector runs in ``"random"`` (or explorative) mode
        exploration_type=ExplorationType.RANDOM,
        # We set the all the devices to be identical. Below is an example of
        # heterogeneous devices
        device=device,
        storing_device=device,
        split_trajs=False,
        postproc=MultiStep(gamma=gamma, n_steps=5),
    )
    return data_collector
```

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

### OpenXExperienceReplay Initialization and Usage Example

Source: https://docs.pytorch.org/rl/0.12/reference/generated/torchrl.data.datasets.OpenXExperienceReplay.html

This snippet demonstrates how to initialize and use the OpenXExperienceReplay dataset with custom transformations like CatTensors and Compose.

```APIDOC
## OpenXExperienceReplay

### Description
Initializes the OpenXExperienceReplay dataset with specified parameters.

### Parameters
- **is_shared** (bool) - Whether the dataset is shared.
- **samplers** (Sampler) - The sampler to use for the dataset.
- **writers** (Writer) - The writer to use for the dataset.
- **batch_size** (int) - The batch size for the dataset.
- **transform** (Transform) - The transformation to apply to the dataset.
- **collate_fn** (callable) - The collate function to use.

### Usage Example
```python
from torchrl.envs import CatTensors, Compose
from tempfile import TemporaryDirectory

cat_tensors = CatTensors(
    in_keys=[("observation", "observation"), ("observation", "achieved_goal"),
             ("observation", "desired_goal")],
    out_key="obs"
    )
cat_next_tensors = CatTensors(
    in_keys=[("next", "observation", "observation"),
             ("next", "observation", "achieved_goal"),
             ("next", "observation", "desired_goal")],
    out_key=("next", "obs")
    )
t = Compose(cat_tensors, cat_next_tensors)

def func(td):
    td = td.select(
        "action",
        "episode",
        ("next", "done"),
        ("next", "observation"),
    )
    return td

# Assuming OpenXExperienceReplay is initialized elsewhere with appropriate parameters
# Example initialization (parameters may vary based on actual dataset structure):
# replay_buffer = OpenXExperienceReplay(batch_size=32, transform=t, collate_fn=func)
```
```

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

### RetrieveKL Transform Example

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

Demonstrates the setup and usage of the RetrieveKL transform for computing KL divergence between two language models. This example requires setting up data structures like History and wrappers for models and policies, along with tokenizer configurations.

```python
>>> from torchrl.data.llm import History
>>> from torchrl.modules.llm import TransformersWrapper
>>> from torchrl.modules.llm.policies import ChatHistory
>>> from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM
>>> from tensordict import TensorDict, set_list_to_stack
>>> import torch
>>>
>>> # Set up list to stack for History
>>> set_list_to_stack(True).set()
>>>
>>> # Create chat data
>>> chats = [
...     [
...         {"role": "system", "content": "You are a helpful assistant."}, 
...         {"role": "user", "content": "Hello, how are you?"},

```

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

### 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. This is useful for accessing nested components.

```python
>>> A.net_b.net_c.conv
Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
```

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

### GymEnv Example

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

Demonstrates how to initialize and inspect the output specification of a Gym environment within TorchRL.

```APIDOC
## GymEnv Initialization and Output Spec

### Description
This example shows how to create a `GymEnv` instance for the "Pendulum-v1" environment and access its `output_spec` attribute to understand the environment's observation, reward, and done specifications.

### Code Example
```python
from torchrl.envs.libs.gym import GymEnv

env = GymEnv("Pendulum-v1")
print(env.output_spec)
```

### Output Specification Example
```
Composite(
    full_reward_spec: Composite(
        reward: UnboundedContinuous(
            shape=torch.Size([1]),
            space=None,
            device=cpu,
            dtype=torch.float32,
            domain=continuous), device=cpu, shape=torch.Size([])),
    full_observation_spec: Composite(
        observation: BoundedContinuous(
            shape=torch.Size([3]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True),
                high=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True)),
            device=cpu,
            dtype=torch.float32,
            domain=continuous), device=cpu, shape=torch.Size([])),
    full_done_spec: Composite(
        done: Categorical(
            shape=torch.Size([1]),
            space=DiscreteBox(n=2),
            device=cpu,
            dtype=torch.bool,
            domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([])
)
```
```

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

### Get Extra State Example

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

Returns any extra state that should be included in the module's state_dict. This is useful for custom state management.

```python
>>> env.get_extra_state()
...
```

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

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