### Install metacontroller-pytorch

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Install the library using pip. No additional setup is required for basic usage.

```bash
pip install metacontroller-pytorch
```

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

### Full Training Pipeline Example

Source: https://context7.com/lucidrains/metacontroller/llms.txt

A complete example demonstrating the three training phases: behavioral cloning, discovery, and internal RL (GRPO). It includes model setup, optimizer initialization, and data sampling.

```python
import torch
from torch.optim import AdamW
from metacontroller import Transformer, MetaController
from metacontroller.metacontroller import extract_grpo_data, z_score

# Model setup
meta_controller = MetaController(
    dim_model=512,
    dim_meta_controller=256,
    dim_latent=128,
    kl_loss_weight=0.2,
    kl_loss_warmup_steps=100,
    ratio_loss_weight=0.1,
    target_temporal_segment_len=4
)

model = Transformer(
    dim=512,
    action_embed_readout=dict(num_discrete=4),
    state_embed_readout=dict(num_continuous=384),
    lower_body=dict(depth=2),
    upper_body=dict(depth=2),
    meta_controller=meta_controller
)

# Optimizers
bc_optimizer = AdamW(model.parameters(), lr=1e-4)
discovery_optimizer = AdamW(meta_controller.discovery_parameters(), lr=1e-4)
rl_optimizer = AdamW(meta_controller.internal_rl_parameters(), lr=1e-4)

# Sample data
state = torch.randn(4, 64, 384)
actions = torch.randint(0, 4, (4, 64))

# === PHASE 1: Behavioral Cloning ===
print("=== Phase 1: Behavioral Cloning ===")
model.train()
for epoch in range(3):
    bc_optimizer.zero_grad()
    state_loss, action_loss = model(state, actions)
    loss = state_loss + action_loss
    loss.backward()
    bc_optimizer.step()
    print(f"BC Epoch {epoch}: loss={loss.item():.4f}")

# === PHASE 2: Discovery Phase ===
print("\n=== Phase 2: Discovery Phase ===")
model.train_discovery()
model.meta_controller_reset_kl_loss_warmup()

for epoch in range(5):
    discovery_optimizer.zero_grad()
    _, action_recon, kl_loss, ratio_loss = model(
        state, actions, discovery_phase=True
    )
    loss = action_recon + kl_loss + ratio_loss
    loss.backward()
    discovery_optimizer.step()
    model.meta_controller_maybe_increment_kl_loss_step()
    print(f"Discovery Epoch {epoch}: recon={action_recon.item():.4f}, kl={kl_loss.item():.4f}")

# === PHASE 3: Internal RL (GRPO) ===
print("\n=== Phase 3: Internal RL ===")
model.train_internal_rl()

# Collect trajectories
for epoch in range(3):
    all_episodes = []
    all_rewards = []

    for group_idx in range(4):
        cache = None
        past_action = None
        grpo_data_list = []

        single_state = state[group_idx:group_idx+1]
        for t in range(single_state.shape[1]):
            t_state = single_state[:, t:t+1]
            with torch.no_grad():
                logits, cache = model(
                    t_state, past_action,
                    meta_controller=meta_controller,
                    return_cache=True
                )
            past_action = model.action_readout.sample(logits)
            grpo_data_list.append(extract_grpo_data(meta_controller, cache))

        states, log_probs, switch_betas, latent_actions = zip(*grpo_data_list)
        all_episodes.append(
            torch.cat(states, dim=1),
            torch.cat(log_probs, dim=1),
            torch.cat(switch_betas, dim=1),
            torch.cat(latent_actions, dim=1)
        )
        all_rewards.append(torch.randn(1))  # Simulated reward

    # Compute advantages and update policy
    advantages = z_score(torch.cat(all_rewards))
    group_states = torch.cat([e[0] for e in all_episodes])
    group_log_probs = torch.cat([e[1] for e in all_episodes])
    group_switch_betas = torch.cat([e[2] for e in all_episodes])
    group_latent_actions = torch.cat([e[3] for e in all_episodes])

    rl_optimizer.zero_grad()
    rl_loss = meta_controller.policy_loss(
        group_states, group_log_probs, group_latent_actions,
        advantages, group_switch_betas > 0.5
    )
    rl_loss.backward()
    rl_optimizer.step()
    print(f"RL Epoch {epoch}: policy_loss={rl_loss.item():.4f}")

print("\nTraining complete!")
```

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

### Install Dependencies and Run Tests

Source: https://github.com/lucidrains/metacontroller/blob/main/README.md

Commands for setting up the development environment and executing the test suite.

```shell
$ uv sync --extra test
```

```shell
$ uv run pytest
```

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

### Install MetaController

Source: https://github.com/lucidrains/metacontroller/blob/main/README.md

Install the package via pip.

```shell
$ pip install metacontroller-pytorch
```

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

### Configure KL Loss Warmup for Discovery Phase

Source: https://context7.com/lucidrains/metacontroller/llms.txt

This example shows how to configure the `MetaController` to apply KL loss warmup. It involves setting the target KL loss weight, the number of warmup steps, and enabling the KL loss weight application.

```python
import torch
from metacontroller import Transformer, MetaController

# Configure KL warmup
meta_controller = MetaController(
    dim_model=512,
    dim_meta_controller=256,
    dim_latent=128,
    kl_loss_weight=0.2,
    kl_loss_warmup_steps=2000,  # Warmup over 2K steps
    apply_kl_loss_weight=True
)

model = Transformer(
    dim=512,
    action_embed_readout=dict(num_discrete=4),
    state_embed_readout=dict(num_continuous=384),
    lower_body=dict(depth=2),
    upper_body=dict(depth=2),
    meta_controller=meta_controller
)

state = torch.randn(2, 32, 384)
actions = torch.randint(0, 4, (2, 32))
```

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

### Initialize Binary Mapper MetaController and Transformer

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Initializes a MetaController with a binary mapper and a Transformer model. This setup is used for tasks involving discrete codes instead of continuous latents.

```python
from metacontroller import Transformer, MetaControllerWithBinaryMapper

meta_controller = MetaControllerWithBinaryMapper(
    dim_model=512,
    dim_meta_controller=256,
    dim_code_bits=4,  # 2^4 = 16 discrete codes
    kl_loss_threshold=0.1,
    switch_temperature=0.1,
    hypernetwork_low_rank=8,
    target_temporal_segment_len=4,
    ratio_loss_weight=0.1,
    kl_loss_weight=0.2,
    kl_loss_warmup_steps=2000
)

model = Transformer(
    dim=512,
    action_embed_readout=dict(num_discrete=4),
    state_embed_readout=dict(num_continuous=384),
    lower_body=dict(depth=2),
    upper_body=dict(depth=2),
    meta_controller=meta_controller
)
```

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

### Initialize MetaController and Transformer for Evolutionary Strategies

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Initializes a MetaController and Transformer model for use with evolutionary strategies. This setup allows for evolving the meta controller's internal RL parameters.

```python
import torch
from metacontroller import Transformer, MetaController

meta_controller = MetaController(
    dim_model=512,
    dim_meta_controller=256,
    dim_latent=128
)

model = Transformer(
    dim=512,
    action_embed_readout=dict(num_discrete=4),
    state_embed_readout=dict(num_continuous=384),
    lower_body=dict(depth=2),
    upper_body=dict(depth=2),
    meta_controller=meta_controller
)
```

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

### Initialize MetaController and Transformer for GRPO

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Initializes a standard MetaController and a Transformer model. This setup is used for collecting trajectories and performing policy optimization in the internal RL phase.

```python
import torch
from metacontroller import Transformer, MetaController

# Initialize model with meta controller
meta_controller = MetaController(
    dim_model=512,
    dim_meta_controller=256,
    dim_latent=128
)

model = Transformer(
    dim=512,
    action_embed_readout=dict(num_discrete=4),
    state_embed_readout=dict(num_continuous=384),
    lower_body=dict(depth=2),
    upper_body=dict(depth=2),
    meta_controller=meta_controller
)
```

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

### Visualize Start and End Frames

Source: https://github.com/lucidrains/metacontroller/blob/main/visualize_trajectories.ipynb

Uses Matplotlib to display the first and last frames of multiple episodes for comparison.

```python
# --- First & last frame comparison for several episodes ---
n_show = min(8, len(lens))
fig, axes = plt.subplots(2, n_show, figsize=(2.2 * n_show, 5))
if n_show == 1:
    axes = axes.reshape(2, 1)

for col in range(n_show):
    ep_len = int(lens[col])
    axes[0, col].imshow(render_frame(states[col, 0]))
    axes[0, col].set_title(f"ep {col} t=0", fontsize=8)
    axes[0, col].axis("off")
    axes[1, col].imshow(render_frame(states[col, ep_len - 1]))
    axes[1, col].set_title(f"ep {col} t={ep_len-1}", fontsize=8)
    axes[1, col].axis("off")

axes[0, 0].set_ylabel("Start", fontsize=10)
axes[1, 0].set_ylabel("End", fontsize=10)
fig.suptitle("Start vs End frames", fontsize=12)
plt.tight_layout()
plt.show()
```

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

### Initialize and Train MetaController

Source: https://github.com/lucidrains/metacontroller/blob/main/README.md

Demonstrates model initialization, behavioral cloning, discovery phase, and internal RL using GRPO.

```python
import torch
from metacontroller import Transformer, MetaController

# 1. initialize model

model = Transformer(
    dim = 512,
    action_embed_readout = dict(num_discrete = 4),
    state_embed_readout = dict(num_continuous = 384),
    lower_body = dict(depth = 2),
    upper_body = dict(depth = 2)
)

state = torch.randn(2, 128, 384)
actions = torch.randint(0, 4, (2, 128))

# 2. behavioral cloning (BC)

state_loss, action_loss = model(state, actions)
(state_loss + action_loss).backward()

# 3. discovery phase

meta_controller = MetaController(
    dim_model = 512,
    dim_meta_controller = 256,
    dim_latent = 128
)

state_pred_loss, action_recon_loss, kl_loss, aux_ratio_loss = model(
    state,
    actions,
    meta_controller = meta_controller,
    discovery_phase = True
)

# they did not use state pred loss in the paper (weight set to 0, but available)
# the ratio loss from h-net paper is also available, but optional (set ratio_loss_weight > 0)

(action_recon_loss + kl_loss * 0.1).backward()

# 4. internal rl phase (GRPO)

# ... collect trajectories ...

logits, cache = model(
    one_state,
    past_action_id,
    meta_controller = meta_controller,
    return_cache = True
)

meta_output = cache.prev_hiddens.meta_controller
old_log_probs = meta_controller.log_prob(meta_output.action_dist, meta_output.actions)

# ... calculate advantages ...

# for GRPO, the inputs to policy loss should be of shape (batch, seq, dim_latent)
# where dim_latent is the dimension of the latent action space

loss = meta_controller.policy_loss(
    group_states,
    group_old_log_probs,
    group_latent_actions,
    group_advantages,
    group_switch_betas
)

loss.backward()
```

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

### Initialize MetaControllerWithBinaryMapper

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Initialize an alternative MetaController variant that uses binary mapping for discrete latent action spaces, aiming for more stable training than VAE-based approaches. This is a specific implementation for discrete latent actions.

```python
import torch
from metacontroller.metacontroller_with_binary_mapper import MetaControllerWithBinaryMapper
from metacontroller import Transformer
```

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

### Training Loop with KL Warmup

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Illustrates a training loop with KL warmup for the Metacontroller. It shows how to increment the KL warmup step counter and reset it when necessary. The current KL loss weight is printed at intervals.

```python
for step in range(100):
    losses, meta_output = model(
        state, actions,
        discovery_phase=True,
        return_meta_controller_output=True
    )

    _, action_recon, kl_loss, ratio_loss = losses
    total_loss = action_recon + kl_loss + ratio_loss
    total_loss.backward()

    # Increment KL warmup step counter
    model.meta_controller_maybe_increment_kl_loss_step()

    if step % 25 == 0:
        current_weight = model.meta_controller_current_kl_loss_weight
        print(f"Step {step}: KL weight = {current_weight:.4f}")

# Reset warmup if needed (e.g., starting new discovery phase)
model.meta_controller_reset_kl_loss_warmup()
print(f"After reset: KL weight = {model.meta_controller_current_kl_loss_weight:.4f}")
```

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

### Load Replay Buffer Data

Source: https://github.com/lucidrains/metacontroller/blob/main/visualize_trajectories.ipynb

Initializes a ReplayBuffer from a directory and prints metadata about the stored trajectories.

```python
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from pathlib import Path
from memmap_replay_buffer import ReplayBuffer

TRAJ_DIR = Path("miniboss-easy-100seeds-100K")

ACTION_NAMES = {
    0: "left",
    1: "right",
    2: "forward",
    3: "pickup",
    4: "drop",
    5: "toggle",
    6: "done",
}

rb = ReplayBuffer.from_folder(TRAJ_DIR)
seeds = np.load(TRAJ_DIR / "seeds.npy")

print(f"Trajectory dir : {TRAJ_DIR}")
print(f"Num episodes   : {rb.num_episodes}")
print(f"Max episodes   : {rb.max_episodes}")
print(f"Max timesteps  : {rb.max_timesteps}")
print(f"State shape    : {rb.shapes['state']}")
print(f"Seeds          : {seeds}")
```

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

### Initialize and Use MetaController for Discovery Phase

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Initialize the MetaController module for learning temporal abstractions. Integrate it into the Transformer model and use for the discovery phase to compute state prediction, action reconstruction, KL, and ratio losses. KL loss has a warmup period.

```python
import torch
from metacontroller import MetaController, Transformer

# Initialize MetaController with configurable parameters
meta_controller = MetaController(
    dim_model=512,
    dim_meta_controller=256,
    dim_latent=128,
    switch_temperature=1.0,
    target_temporal_segment_len=4,
    ratio_loss_weight=0.1,
    kl_loss_weight=0.2,
    kl_loss_warmup_steps=2000,  # Linear KL warmup over 2K steps
    internal_sequence_embedder=dict(
        attn_dim_head=32,
        heads=8,
        depth=2,
        polar_pos_emb=True
    ),
    action_proposer=dict(
        depth=2,
        attn_dim_head=32,
        heads=8,
        polar_pos_emb=True
    )
)

# Initialize transformer with meta controller
model = Transformer(
    dim=512,
    action_embed_readout=dict(num_discrete=4),
    state_embed_readout=dict(num_continuous=384),
    lower_body=dict(depth=2),
    upper_body=dict(depth=2),
    meta_controller=meta_controller
)

state = torch.randn(2, 128, 384)
actions = torch.randint(0, 4, (2, 128))

# Phase 2: Discovery Phase - learns temporal abstractions
state_pred_loss, action_recon_loss, kl_loss, ratio_loss = model(
    state,
    actions,
    discovery_phase=True
)

# Total discovery loss (state pred weight is 0 in paper)
discovery_loss = action_recon_loss + kl_loss + ratio_loss
discovery_loss.backward()

print(f"Discovery - Action Recon: {action_recon_loss.item():.4f}")
print(f"Discovery - KL Loss: {kl_loss.item():.4f}")
print(f"Discovery - Ratio Loss: {ratio_loss.item():.4f}")
```

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

### Run Evolutionary Strategies with GRPO

Source: https://context7.com/lucidrains/metacontroller/llms.txt

This snippet shows how to initiate evolutionary strategies over GRPO using the `evolve` method. It requires specifying the number of generations, the environment callable, and the noise population size.

```python
model.evolve(
    num_generations=10,
    environment=environment_callable,
    noise_population_size=4
)
print("Evolution complete!")
```

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

### Initialize TransformerForSymbolicGrid for BabyAI

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Demonstrates initializing `TransformerForSymbolicGrid` for symbolic grid environments like BabyAI. It requires specifying dimensions, number of discrete channels, grid size, and architectural details for spatial processing and action/state embeddings.

```python
import torch
from metacontroller import MetaController
from metacontroller.transformer_for_symbolic_grid import TransformerForSymbolicGrid

# For BabyAI-style symbolic observations
model = TransformerForSymbolicGrid(
    dim=512,
    num_discrete=(11, 6, 3),  # Type, Color, State channels
    grid_size=7,  # 7x7 grid
    spatial_processor_depth=4,
    convnext_kernel_size=3,
    action_embed_readout=dict(num_discrete=7),  # BabyAI has 7 actions
    lower_body=dict(depth=2),
    upper_body=dict(depth=2)
)

# Symbolic state: (batch, time, width, height, channels)
batch_size = 2
seq_len = 64
symbolic_state = torch.randint(0, 11, (batch_size, seq_len, 7, 7, 3))
actions = torch.randint(0, 7, (batch_size, seq_len))

# Behavioral cloning
state_loss, action_loss = model(symbolic_state, actions)
total_loss = state_loss + action_loss
total_loss.backward()

print(f"Symbolic Grid - State Loss: {state_loss.item():.4f}")
print(f"Symbolic Grid - Action Loss: {action_loss.item():.4f}")
```

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

### Initialize and Use Transformer for Behavioral Cloning

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Initialize the main Transformer model for sequence modeling. Use this for the behavioral cloning phase, providing state and actions to compute losses. Supports variable-length episodes using episode_lens.

```python
import torch
from metacontroller import Transformer, MetaController

# Initialize the transformer model
model = Transformer(
    dim=512,
    action_embed_readout=dict(num_discrete=4),  # 4 discrete actions
    state_embed_readout=dict(num_continuous=384),  # continuous state space
    lower_body=dict(depth=2, heads=8, attn_dim_head=64),
    upper_body=dict(depth=2, heads=8, attn_dim_head=64),
    embed_past_actions=True,
    normalize_state_action_losses=False
)

# Sample data
batch_size = 2
seq_len = 128
state = torch.randn(batch_size, seq_len, 384)
actions = torch.randint(0, 4, (batch_size, seq_len))

# Phase 1: Behavioral Cloning
model.train()
state_loss, action_loss = model(state, actions)
total_loss = state_loss + action_loss
total_loss.backward()
print(f"BC Loss - State: {state_loss.item():.4f}, Action: {action_loss.item():.4f}")

# For variable-length episodes, use episode_lens
episode_lens = torch.tensor([100, 128])
state_loss, action_loss = model(state, actions, episode_lens=episode_lens)
```

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

### Evolve MetaController with Evolutionary Strategies

Source: https://github.com/lucidrains/metacontroller/blob/main/README.md

Alternative training approach using evolutionary strategies for the final phase.

```python
# 5. evolve (ES over GRPO)

model.meta_controller = meta_controller

def environment_callable(model):
    # return a fitness score
    return 1.0

model.evolve(
    num_generations = 10,
    environment = environment_callable
)
```

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

### Save and Load Transformer and MetaController Models

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Shows how to save and load `Transformer` and `MetaController` models using the `@save_load()` decorator. This involves initializing the models, saving their states, and then loading them back into new instances.

```python
import torch
from metacontroller import Transformer, MetaController

# Initialize and train model
meta_controller = MetaController(
    dim_model=512,
    dim_meta_controller=256,
    dim_latent=128
)

model = Transformer(
    dim=512,
    action_embed_readout=dict(num_discrete=4),
    state_embed_readout=dict(num_continuous=384),
    lower_body=dict(depth=2),
    upper_body=dict(depth=2),
    meta_controller=meta_controller
)

# Save models
model.save('./transformer_trained.pt')
meta_controller.save('./meta_controller_trained.pt')
print("Models saved!")

# Load models (creates new instance with saved weights and config)
loaded_model = Transformer.init_and_load('./transformer_trained.pt', strict=False)
loaded_meta_controller = MetaController.init_and_load('./meta_controller_trained.pt')
print("Models loaded!")

# Verify loaded model works
state = torch.randn(1, 16, 384)
actions = torch.randint(0, 4, (1, 16))

with torch.no_grad():
    state_loss, action_loss = loaded_model(state, actions)
    print(f"Loaded model - State Loss: {state_loss.item():.4f}")
```

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

### Sample and Display Trajectories

Source: https://github.com/lucidrains/metacontroller/blob/main/visualize_trajectories.ipynb

Uses numpy to select random indices from a dataset and displays the corresponding trajectories.

```python
rng = np.random.default_rng(42)
sample_ids = rng.choice(len(lens), size=min(5, len(lens)), replace=False)

for idx in sample_ids:
    print(f"max: {states[idx].max()}")
    print(f"tensor: {states[idx]}")
    show_trajectory(states[idx], actions[idx], int(lens[idx]), ep_idx=idx)
```

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

### Ablate Switch Beta with Regular Switching Pattern

Source: https://context7.com/lucidrains/metacontroller/llms.txt

This snippet demonstrates how to create and apply a regular switching pattern for the `switch_beta` parameter in the `Transformer` model for ablation studies. It shows how to generate the pattern and use it during model inference.

```python
import torch
from metacontroller import Transformer, MetaController

meta_controller = MetaController(
    dim_model=512,
    dim_meta_controller=256,
    dim_latent=128
)

model = Transformer(
    dim=512,
    action_embed_readout=dict(num_discrete=4),
    state_embed_readout=dict(num_continuous=384),
    lower_body=dict(depth=2),
    upper_body=dict(depth=2),
    meta_controller=meta_controller
)

batch_size = 2
seq_len = 16
state = torch.randn(batch_size, seq_len, 384)
actions = torch.randint(0, 4, (batch_size, seq_len))

# Create regular switching pattern (switch every N steps)
frequency = 4
ablate_switch_beta = MetaController.create_regular_switch_beta(
    batch=batch_size,
    seq_len=seq_len,
    frequency=frequency
)
# Result: switches at timesteps 3, 7, 11, 15 (every 4th step)
print(f"Regular switch pattern: {ablate_switch_beta[0]}")

# Discovery with custom switch pattern
losses, meta_output = model(
    state,
    actions,
    discovery_phase=True,
    ablate_switch_beta=ablate_switch_beta,
    return_meta_controller_output=True
)

# Or use switch_beta_frequency shorthand
losses, meta_output = model(
    state,
    actions,
    discovery_phase=True,
    switch_beta_frequency=4,  # Equivalent to ablate_switch_beta above
    return_meta_controller_output=True
)

print(f"Used switch beta: {meta_output.switch_beta[0]}")
```

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

### Discovery Phase with Binary Mapper

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Performs a discovery phase using the Transformer model with a binary mapper. It calculates losses including action reconstruction and KL divergence.

```python
import torch

state = torch.randn(2, 128, 384)
actions = torch.randint(0, 4, (2, 128))

# Discovery phase with binary mapper
losses = model(state, actions, discovery_phase=True)
state_pred, action_recon, kl_loss, ratio_loss = losses

print(f"Binary Mapper - Num Codes: {meta_controller.num_codes}")
print(f"Binary Mapper - Action Recon: {action_recon.item():.4f}")
```

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

### Visualize Action Distribution

Source: https://github.com/lucidrains/metacontroller/blob/main/visualize_trajectories.ipynb

Calculates and plots the frequency of each action across all loaded episodes.

```python
# --- Action distribution ---
all_actions = []
for i in range(len(lens)):
    all_actions.append(actions[i, :lens[i]])
all_actions = np.concatenate(all_actions).astype(int)

unique_acts, counts = np.unique(all_actions, return_counts=True)
labels = [ACTION_NAMES.get(a, str(a)) for a in unique_acts]

fig, ax = plt.subplots(figsize=(6, 3))
bars = ax.bar(labels, counts, color="steelblue", edgecolor="white", linewidth=0.5)
for bar, c in zip(bars, counts):
    ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height(),
            f"{c/len(all_actions)*100:.1f}%", ha="center", va="bottom", fontsize=8)
ax.set_xlabel("Action")
ax.set_ylabel("Count")
ax.set_title("Action Distribution (all episodes)")
plt.tight_layout()
plt.show()
```

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

### Compute Batch Statistics

Source: https://github.com/lucidrains/metacontroller/blob/main/visualize_trajectories.ipynb

Loads a batch of episodes from the replay buffer to calculate and print range and length statistics.

```python
# Load a large batch to compute statistics
STAT_BATCH = min(rb.num_episodes, 4096)
dl = rb.dataloader(batch_size=STAT_BATCH)
batch = next(iter(dl))

states = batch["state"].numpy()        # (B, T, 7, 7, 3)
actions = batch["action"].numpy()      # (B, T)
lens = batch["_lens"].numpy()          # (B,)
mission_emb = batch["mission_embedding"].numpy()  # (B, 384)

print(f"Loaded batch: {STAT_BATCH} episodes")
print(f"State range  : [{states.min():.0f}, {states.max():.0f}]")
print(f"Action range : [{actions.min():.0f}, {actions.max():.0f}]")
print(f"Ep lengths   : min={lens.min()}, max={lens.max()}, mean={lens.mean():.1f}, median={np.median(lens):.1f}")
```

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

### Generate Action Sequence Heatmap

Source: https://github.com/lucidrains/metacontroller/blob/main/visualize_trajectories.ipynb

Visualizes action sequences for the first N episodes using a heatmap. Requires matplotlib and numpy. Note the MatplotlibDeprecationWarning for get_cmap.

```python
N_HEAT = min(40, len(lens))
max_t = int(lens[:N_HEAT].max())
action_grid = np.full((N_HEAT, max_t), np.nan)
for i in range(N_HEAT):
    l = int(lens[i])
    action_grid[i, :l] = actions[i, :l]

fig, ax = plt.subplots(figsize=(12, 5))
cmap = plt.cm.get_cmap("tab10", len(ACTION_NAMES))
im = ax.imshow(action_grid, aspect="auto", cmap=cmap, vmin=-0.5, vmax=len(ACTION_NAMES) - 0.5, interpolation="nearest")
cbar = fig.colorbar(im, ax=ax, ticks=list(ACTION_NAMES.keys()))
cbar.ax.set_yticklabels([ACTION_NAMES[k] for k in ACTION_NAMES])
ax.set_xlabel("Timestep")
ax.set_ylabel("Episode")
ax.set_title("Action Sequences (first episodes)")
plt.tight_layout()
plt.show()
```

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

### Simulate Trajectory Collection and Compute Policy Loss

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Simulates collecting multiple trajectories (episodes) and computes the policy loss using GRPO. It involves z-scoring advantages and preparing data for the policy loss function.

```python
# Simulate trajectory collection for multiple episodes (GRPO groups)
state = torch.randn(1, 16, 384)
all_rewards = []
all_episodes = []

for _ in range(3):  # Group of 3 trajectories
    grpo_data_list = collect_trajectory(model, meta_controller, state)

    states, actions, log_probs, switch_betas = zip(*grpo_data_list)
    all_episodes.append((
        torch.cat(states, dim=1),
        torch.cat(log_probs, dim=1),
        torch.cat(switch_betas, dim=1),
        torch.cat(actions, dim=1)
    ))
    all_rewards.append(torch.randn(1))  # Simulated reward

# Compute z-scored advantages
rewards = torch.cat(all_rewards)
group_advantages = z_score(rewards)

# Combine episode data
group_states = torch.cat([ep[0] for ep in all_episodes], dim=0)
group_log_probs = torch.cat([ep[1] for ep in all_episodes], dim=0)
group_switch_betas = torch.cat([ep[2] for ep in all_episodes], dim=0)
group_latent_actions = torch.cat([ep[3] for ep in all_episodes], dim=0)

# Phase 3: Internal RL Policy Update
meta_controller.train_internal_rl()

loss = meta_controller.policy_loss(
    group_states,
    group_log_probs,
    group_latent_actions,
    group_advantages,
    group_switch_betas > 0.5,  # Use hard switch mask
    eps_clip=0.2
)

loss.backward()
print(f"Policy Loss: {loss.item():.4f}")
```

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

### Render Trajectory Frames

Source: https://github.com/lucidrains/metacontroller/blob/main/visualize_trajectories.ipynb

Provides utility functions to upscale frames and display a sequence of states from a trajectory.

```python
def render_frame(frame, scale=12):
    """Upscale a tiny (7,7,3) frame for display via nearest-neighbor."""
    img = frame.astype(np.uint8)
    return np.kron(img, np.ones((scale, scale, 1))).astype(np.uint8)


def show_trajectory(states, actions, ep_len, ep_idx, max_frames=20, scale=12):
    """Display frames from a single trajectory in a grid."""
    n = min(ep_len, max_frames)
    step = max(1, ep_len // n)
    indices = list(range(0, ep_len, step))[:n]

    fig, axes = plt.subplots(1, len(indices), figsize=(2 * len(indices), 2.5))
    if len(indices) == 1:
        axes = [axes]
    for ax, t in zip(axes, indices):
        ax.imshow(render_frame(states[t], scale=scale))
        act = int(actions[t])
        ax.set_title(f"t={t}\n{ACTION_NAMES.get(act, act)}", fontsize=8)
        ax.axis("off")
    fig.suptitle(f"Episode {ep_idx}  (len={ep_len})", fontsize=11)
    plt.tight_layout()
    plt.show()
```

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

### Define Environment Fitness Function for Evolutionary Strategies

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Defines a callable function that simulates running a model in an environment and returns a fitness score. This function is used by evolutionary strategies to evaluate model performance.

```python
# Define environment fitness function
def environment_callable(model):
    """
    Run the model in an environment and return fitness score.
    Higher is better.
    """
    # Simulate environment interaction
    state = torch.randn(1, 16, 384)
    with torch.no_grad():
        logits, _ = model(state, return_cache=True)

    # Return simulated fitness (replace with actual env reward)
    return torch.randn(1).item()
```

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

### Visualize Mission Embeddings with PCA

Source: https://github.com/lucidrains/metacontroller/blob/main/visualize_trajectories.ipynb

Performs Principal Component Analysis (PCA) on mission embeddings to visualize them in 2D, colored by episode length. Requires scikit-learn and matplotlib.

```python
# --- Mission embedding PCA (2D) ---
from sklearn.decomposition import PCA

pca = PCA(n_components=2)
emb_2d = pca.fit_transform(mission_emb)

fig, ax = plt.subplots(figsize=(6, 5))
sc = ax.scatter(emb_2d[:, 0], emb_2d[:, 1], c=lens, cmap="viridis", s=8, alpha=0.6)
fig.colorbar(sc, ax=ax, label="Episode length")
ax.set_xlabel(f"PC1 ({pca.explained_variance_ratio_[0]*100:.1f}%)")
ax.set_ylabel(f"PC2 ({pca.explained_variance_ratio_[1]*100:.1f}%)")
ax.set_title("Mission Embeddings (PCA) coloured by episode length")
plt.tight_layout()
plt.show()
```

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

### Visualize Episode Length Distribution

Source: https://github.com/lucidrains/metacontroller/blob/main/visualize_trajectories.ipynb

Generates a histogram showing the distribution of episode lengths using matplotlib.

```python
# --- Episode length distribution ---
fig, ax = plt.subplots(figsize=(8, 3))
ax.hist(lens, bins=50, edgecolor="white", linewidth=0.5, color="steelblue")
ax.axvline(lens.mean(), color="tomato", ls="--", label=f"mean={lens.mean():.1f}")
ax.axvline(np.median(lens), color="orange", ls="--", label=f"median={np.median(lens):.1f}")
ax.set_xlabel("Episode length (timesteps)")
ax.set_ylabel("Count")
ax.set_title("Episode Length Distribution")
ax.legend()
plt.tight_layout()
plt.show()
```

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

### Collect Trajectory Data for GRPO

Source: https://context7.com/lucidrains/metacontroller/llms.txt

Defines a function to collect trajectory data, including GRPO-specific information like logits, cache, and past action IDs, during environment interaction.

```python
from metacontroller.metacontroller import extract_grpo_data, z_score

# Collect trajectories iteratively during environment interaction
def collect_trajectory(model, meta_controller, initial_state):
    cache = None
    past_action_id = None
    grpo_data_list = []

    for timestep_state in initial_state.unbind(dim=1):
        timestep_state = timestep_state.unsqueeze(1)  # (batch, 1, dim)

        with torch.no_grad():
            logits, cache = model(
                timestep_state,
                past_action_id,
                meta_controller=meta_controller,
                return_cache=True
            )

        past_action_id = model.action_readout.sample(logits)

        # Extract GRPO data for policy optimization
        grpo_data = extract_grpo_data(meta_controller, cache)
        grpo_data_list.append(grpo_data)

    return grpo_data_list
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.