### Full Flow Steering Training Setup

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/flow-steering-api.md

This example demonstrates the complete setup for training a Flow Steering model. It includes initializing the base model, the RL model with specific hyperparameters like `expectile_tau`, and configuring the optimizer.

```python
import torch
from torch.utils.data import DataLoader
from mimic_video import MimicVideo
from mimic_video.cosmos_predict import CosmosPredictWrapper
from mimic_video.flow_steering import FlowSteering

# 1. Setup models
video_wrapper = CosmosPredictWrapper(extract_layer=19)
base_model = MimicVideo(dim=512, video_predict_wrapper=video_wrapper)
rl_model = FlowSteering(
    base_model,
    discount_factor=0.99,
    ema_decay=0.999,
    expectile_tau=0.25  # Pessimistic for offline RL
)

# 2. Setup optimizer (only train RL components, not video encoder)
critic_params = list(rl_model.outer_critic.parameters()) + list(rl_model.inner_critic.parameters())
actor_params = list(rl_model.actor.parameters())
optimizer = torch.optim.AdamW(
    list(critic_params) + list(actor_params),
    lr=1e-4
)

```

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

### Multi-View Setup and Training Example

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/configuration.md

Configures Mimic Video for multi-view input (e.g., from multiple cameras). Demonstrates how to pass multi-view video data during training.

```python
model = MimicVideo(
    dim=512,
    video_predict_wrapper=wrapper,
    num_video_viewpoints=3,  # 3 cameras
    extracted_video_layer_indices=[19]*8
)

# Training: pass (batch, num_views, frames, channels, height, width)
loss = model(
    prompts=['task'],
    video=torch.rand(2, 3, 5, 3, 256, 256),  # batch=2, views=3
    joint_state=torch.randn(2, 32),
    actions=torch.randn(2, 32, 20)
)
```

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

### Basic Supervised Learning Example

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/INDEX.txt

Demonstrates the fundamental setup for supervised learning using the MimicVideo class. Requires importing the necessary class.

```python
from mimic_video.mimic_video import MimicVideo

model = MimicVideo(
    num_tokens = 8192,
    dim = 512,
    depth = 6,
    heads = 8
)

# Dummy data
video = torch.randn(1, 3, 128, 128)

# Forward pass
loss = model(video)

# Backward pass and optimization
loss.backward()
optimizer.step()

```

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

### Minimal Mimic-Video Example

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/README.md

A minimal example demonstrating the setup, training, and inference of the MimicVideo model. Requires torch and mimic-video libraries.

```python
import torch
from mimic_video import MimicVideo
from mimic_video.cosmos_predict import CosmosPredictWrapper

# 1. Setup video encoder
wrapper = CosmosPredictWrapper(
    extract_layer=19,
    random_weights=True,  # For testing
    tiny=True
)

# 2. Create action model
model = MimicVideo(512, wrapper)

# 3. Training
video = torch.rand(2, 5, 3, 32, 32)
joint_state = torch.randn(2, 32)
actions = torch.randn(2, 32, 20)

loss = model(
    prompts=['task 1', 'task 2'],
    video=video,
    joint_state=joint_state,
    actions=actions
)
loss.backward()

# 4. Inference
sampled_actions = model.sample(
    prompts='new task',
    video=video[:1],
    joint_state=joint_state[:1]
)
```

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

### MimicVideo Model Initialization and Training Example

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/mimic-video-api.md

Demonstrates how to initialize the MimicVideo model with a video encoder wrapper and perform a training forward pass. This snippet shows the setup for using pretrained weights and passing various inputs like prompts, video, joint state, and actions to compute the loss.

```python
import torch
from mimic_video import MimicVideo
from mimic_video.cosmos_predict import CosmosPredictWrapper

# Setup video encoder wrapper
video_wrapper = CosmosPredictWrapper(
    model_name='nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
    extract_layer=19,
    random_weights=False  # Use pretrained weights
)

# Create MimicVideo model
model = MimicVideo(
    dim=512,
    video_predict_wrapper=video_wrapper,
    depth=8,
    action_chunk_len=32,
    dim_action=20,
    dim_joint_state=32,
)

# Training: forward pass returns scalar loss
video = torch.rand(2, 5, 3, 256, 256)  # batch_size=2, frames=5
joint_state = torch.randn(2, 32)
actions = torch.randn(2, 32, 20)

loss = model(
    prompts=['put the package on the conveyer belt', 'pass the butter'],
    video=video,
    joint_state=joint_state,
    actions=actions
)

loss.backward()

```

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

### Multi-view Learning Setup

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/INDEX.txt

Illustrates how to configure MimicVideo for multi-view learning scenarios. This setup involves specific parameter adjustments for handling multiple input views.

```python
from mimic_video.mimic_video import MimicVideo

model = MimicVideo(
    num_tokens = 8192,
    dim = 512,
    depth = 6,
    heads = 8,
    # multi_view_learning = True # Example parameter, actual parameter name may vary
)

# Dummy data for multi-view
video_view1 = torch.randn(1, 3, 128, 128)
video_view2 = torch.randn(1, 3, 128, 128)

# Forward pass with multiple views
# loss = model(video_view1, video_view2) # Example call, actual call signature may vary

```

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

### Install mimic-video Package

Source: https://github.com/lucidrains/mimic-video/blob/main/README.md

Install the mimic-video library using pip. This command installs the core package.

```shell
pip install mimic-video
```

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

### Install Test Dependencies for mimic-video

Source: https://github.com/lucidrains/mimic-video/blob/main/README.md

Installs the mimic-video package along with its test dependencies. This is required before running tests.

```shell
pip install '.[test]'
```

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

### CosmosPredictWrapper Constructor Example

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/INDEX.txt

Demonstrates the initialization of the CosmosPredictWrapper class with key parameters. This wrapper is used for video encoding and feature extraction.

```python
from mimic_video.cosmos.predict import CosmosPredictWrapper

model = CosmosPredictWrapper(
    num_tokens = 8192,  # Number of tokens in the vocabulary
    dim = 512,         # Dimensionality of the model
    depth = 6,         # Number of transformer layers
    heads = 8          # Number of attention heads
    # ... other parameters
)

```

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

### Instantiate MimicVideo Model

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/mimic-video-api.md

Instantiate the MimicVideo model with specified dimensions and state autoencoder configuration. This is the initial setup for using the model.

```python
model = MimicVideo(
    dim=512,
    video_wrapper=wrapper,
    state_autoencoder=dict(enc_depth=2, dec_depth=2, max_seq_len=64, dim_latent=32)
)
```

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

### MimicVideo Initialization with Normalization

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/types.md

Example of initializing MimicVideo with mean and standard deviation for automatic action normalization. The model normalizes actions on forward() and denormalizes on sample().

```python
model = MimicVideo(
    dim=512,
    action_mean_std=torch.tensor([
        [0.0, 0.0, ...],  # Mean per action dimension
        [1.0, 1.0, ...]   # Std per action dimension
    ])
)
# Automatically normalizes actions on forward(), denormalizes on sample()
```

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

### FlowSteering Constructor Example

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/INDEX.txt

Demonstrates the initialization of the FlowSteering class with its core parameters. This class is designed for offline reinforcement learning tasks.

```python
from mimic_video.flow_steering.flow_steering import FlowSteering

model = FlowSteering(
    num_tokens = 8192,
    dim = 512,
    depth = 6,
    heads = 8,
    # Additional parameters specific to FlowSteering, e.g., discount factor, etc.
)

```

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

### FlowSteering Usage Example

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/flow-steering-api.md

Demonstrates the complete workflow of using the FlowSteering API, from training a base flow model to wrapping it with RL for offline learning and sampling with exploration.

```python
import torch
from mimic_video import MimicVideo
from mimic_video.cosmos_predict import CosmosPredictWrapper
from mimic_video.flow_steering import FlowSteering

# 1. Train base flow model on behavior data
video_wrapper = CosmosPredictWrapper(extract_layer=19)
flow_model = MimicVideo(dim=512, video_predict_wrapper=video_wrapper)

# Training loop (simplified)
for batch in dataset:
    loss = flow_model(
        prompts=batch.prompts,
        video=batch.video,
        actions=batch.actions,
        joint_state=batch.joint_state
    )
    loss.backward()

# 2. Wrap with RL for offline learning
rl_model = FlowSteering(
    flow_model,
    discount_factor=0.999,
    ema_decay=0.999,
    use_minto=True,
    expectile_tau=0.25  # Pessimistic for offline RL
)

# 3. RL training loop
for batch in rl_dataset:
    loss, (actor_loss, outer_critic_loss, inner_critic_loss) = rl_model(
        prompts=batch.prompts,
        video=batch.video,
        joint_state=batch.joint_state,
        actions=batch.actions,
        noise_latents=batch.noise_latents,
        next_video=batch.next_video,
        next_joint_state=batch.next_joint_state,
        rewards=batch.rewards,
        done=batch.done
    )
    loss.backward()
    optimizer.step()
    
    # Update EMA targets
    rl_model.update_critic_ema()

# 4. Sampling with exploration
actions, noise_latents = rl_model.sample(
    prompts='task',
    video=test_video,
    joint_state=test_joint_state,
    exploration_noise_std=0.1  # Add exploration
)
```

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

### Real-time Chunking (RTC) Setup

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/INDEX.txt

Illustrates the configuration for Real-time Chunking (RTC) within the mimic-video framework. This pattern is designed for processing video streams in manageable segments.

```python
from mimic_video.mimic_video import MimicVideo

model = MimicVideo(
    num_tokens = 8192,
    dim = 512,
    depth = 6,
    heads = 8,
    # rtc_enabled = True # Example parameter, actual parameter name may vary
)

# Dummy data for RTC
video_chunk = torch.randn(1, 3, 128, 128)

# Process a video chunk
# output = model.process_chunk(video_chunk) # Example call, actual call signature may vary

```

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

### Model Forward Pass for Training

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/mimic-video-api.md

Example of using the `forward` method for training. The method returns the loss, which can then be backpropagated.

```python
# Training: forward returns loss
loss = model(
    actions=actions,
    joint_state=joint_state,
    video=video,
    prompts=['task description']
)
loss.backward()
```

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

### Sample Action with Flow Steering

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/flow-steering-api.md

Use `actor_forward` to sample an action. Set `sample_noise_latent=True` to get a stochastic sample, or `False` to retrieve the mean and standard deviation of the action distribution.

```python
noise = rl_model.actor_forward(
    prompts='task',
    video=video,
    joint_state=joint_state,
    sample_noise_latent=True
)
```

```python
mean, std = rl_model.actor_forward(
    prompts='task',
    video=video,
    joint_state=joint_state,
    sample_noise_latent=False
)
```

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

### Actor-Critic RL Architecture Overview

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/modules-architecture.md

Details the roles and training objectives of the Actor, Outer Critic, and Inner Critic components in the RL setup. The Inner Critic guides the actor to stay within the data distribution.

```markdown
Actor (Policy):
├── Input: video, joint_state, prompts
├── Output: (mean, logvar) of noise latent distribution
└── Training: Loss = -E[Q_inner(s, noise)]
             → Maximize value in noise latent space

Outer Critic (Q-function):
├── Input: (state, action)
├── Output: scalar value estimate
└── Training: Uses discounted returns with expectile loss
             Bellman target: R + γ × EMA_Q(s', a')

Inner Critic (Steering Guide):
├── Input: (state, noise_latent)
├── Output: scalar value estimate
└── Training: MSE to EMA_outer_Q(s, action_from_noise)
             → Keeps actor steering within data distribution
```

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

### MimicVideo Configuration: Medium (Recommended)

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Initialize MimicVideo with recommended medium dimensions, depth, and enable real-time chunking for a balance of performance and capability.

```python
MimicVideo(dim=512, depth=8, action_chunk_len=32, train_time_rtc=True)
```

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

### Initialize Mimic Video Model and Wrapper

Source: https://github.com/lucidrains/mimic-video/blob/main/README.md

Demonstrates how to initialize the CosmosPredictWrapper and the MimicVideo model. The wrapper is used for video processing, and the model integrates it for action prediction.

```python
import torch

from mimic_video.cosmos_predict import CosmosPredictWrapper

video_wrapper = CosmosPredictWrapper(
    extract_layer = 1,
    random_weights = True,
    tiny = True
)

from mimic_video import MimicVideo

model = MimicVideo(512, video_wrapper)
```

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

### MimicVideo Configuration: Small (Testing)

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Initialize MimicVideo with small dimensions and fewer layers, suitable for testing and quick experimentation.

```python
MimicVideo(dim=256, depth=4, action_chunk_len=16)
```

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

### Small Model Configuration

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/configuration.md

Sets up a small Mimic Video model for fast iteration. Initializes with random weights and uses a tiny architecture.

```python
from mimic_video import MimicVideo
from mimic_video.cosmos_predict import CosmosPredictWrapper

wrapper = CosmosPredictWrapper(
    extract_layer=19,
    random_weights=True,
    tiny=True
)

model = MimicVideo(
    dim=256,
    video_predict_wrapper=wrapper,
    depth=4,
    action_chunk_len=16,
    dim_action=7,
    dim_joint_state=7
)
```

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

### Get Latent Dimension

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/cosmos-predict-api.md

Access the computed latent dimension of the wrapper after initialization. This value represents the dimensionality of the extracted features.

```python
print(f"Latent dimension: {wrapper.dim_latent}")  # e.g., 4096
```

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

### Initialize MimicVideo for Real-Time Control (RTC)

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Sets up the MimicVideo model for real-time control by enabling action prefixing. This configuration is intended for low-latency predictions. During inference, use `model.sample` with `prefix_action_chunk`.

```python
# Enable action prefixing for low-latency prediction
model = MimicVideo(
    dim=512,
    train_time_rtc=True,
    train_time_rtc_max_delay=8
)
# At inference: model.sample(prefix_action_chunk=partial_actions)
```

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

### Get Model Device

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/mimic-video-api.md

Access the device property of the model to determine if it is running on CPU or GPU. This is inferred from the model's parameters.

```python
@property
def device(self) -> torch.device
```

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

### Get Action Parameters

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/mimic-video-api.md

Returns a set of model parameters excluding video encoder weights. This is useful for selectively fine-tuning only the action prediction backbone.

```python
# Freeze video encoder, train action decoder only
optimizer = torch.optim.AdamW(
    model.action_parameters(),
    lr=1e-4
)
```

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

### Basic Supervised Learning with MimicVideo

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/usage-patterns.md

Train MimicVideo on demonstration data using supervised learning. This involves setting up the model, optimizer, and a standard training loop for loss calculation and backpropagation. Inference can then be performed using the trained model.

```python
import torch
from torch.utils.data import DataLoader
from mimic_video import MimicVideo
from mimic_video.cosmos_predict import CosmosPredictWrapper

# Setup
wrapper = CosmosPredictWrapper(
    model_name='nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
    extract_layer=19
)

model = MimicVideo(
    dim=512,
    video_predict_wrapper=wrapper,
    depth=8,
    action_chunk_len=32,
    dim_action=7,
    dim_joint_state=7
)

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

# Training loop
model.train()
for epoch in range(num_epochs):
    for batch in dataset:
        # Forward pass: loss from flow matching
        loss = model(
            prompts=batch.task_descriptions,
            video=batch.videos,           # (batch, frames, 3, h, w)
            joint_state=batch.joint_states,
            actions=batch.actions
        )
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        
        print(f"Loss: {loss.item():.4f}")

# Inference
model.eval()
with torch.no_grad():
    test_actions = model.sample(
        steps=16,
        prompts='place object in bin',
        video=test_video,
        joint_state=test_joint_state
    )
```

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

### Move Model to Device

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/cosmos-predict-api.md

Get the device (CPU or CUDA) of the model parameters. This is useful for ensuring that input data is on the same device as the model before performing operations.

```python
video = video.to(wrapper.device)
```

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

### RL Model Usage: Single-step Rewards

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/flow-steering-api.md

Example of using the RL model with single-step rewards. Ensure rewards have shape (batch,) and done has shape (batch,).

```python
# Single-step rewards
loss, (a_loss, q_loss, steer_loss) = rl_model(
    prompts='grasp object',
    video=video,
    joint_state=joint_state,
    actions=actions,
    noise_latents=noise_latents,
    next_video=next_video,
    next_joint_state=next_joint_state,
    rewards=torch.tensor([0.5, 0.8]),  # shape: (batch,)
    done=torch.tensor([False, False])
)
```

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

### Initialize Mimic Video for Offline RL

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/configuration.md

Instantiate the MimicVideo model and then wrap it with FlowSteering for offline RL tasks. Adjust parameters like `num_advantage_ids` for return estimation and `discount_factor`, `expectile_tau` for RL behavior.

```python
from mimic_video.flow_steering import FlowSteering

flow_model = MimicVideo(
    dim=512,
    video_predict_wrapper=wrapper,
    depth=8,
    is_flow_model=True,
    num_advantage_ids=3  # For return estimation
)

rl_model = FlowSteering(
    flow_model,
    discount_factor=0.99,
    ema_decay=0.995,
    expectile_tau=0.25,  # Pessimistic for offline
    use_minto=True
)
```

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

### Apply Gradient Clipping

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Applies gradient clipping to the model's parameters to prevent exploding gradients during training. A maximum norm of 1.0 is used in this example.

```python
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
```

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

### Initialize MimicVideo for Supervised Learning

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Initializes the MimicVideo model for supervised learning tasks. No RL-specific parameters are required for this use case. The model is then used to compute a loss based on input prompts, video, joint state, and actions.

```python
# No RL-specific parameters needed
model = MimicVideo(dim=512, ...)
loss = model(prompts=..., video=..., joint_state=..., actions=...)
```

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

### Verify Model Initialization and Configuration

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/errors.md

Checks the initialization parameters of the MimicVideo model and its components. Useful for confirming dimensions and shapes match expectations.

```python
# Verify initialization
model = MimicVideo(dim=512, dim_video_hidden=4096)
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
print(f"Action shape: {model.action_shape}")
print(f"Video hidden dim: {model.dim_video_hidden}")

# Check video wrapper
wrapper = model.video_predict_wrapper
print(f"Latent dimension: {wrapper.dim_latent if wrapper else 'N/A'}")
```

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

### Initialize CosmosPredictWrapper with LoRA Adapter

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/cosmos-predict-api.md

Instantiate the wrapper with a specified layer extraction and a path to LoRA adapter weights for parameter-efficient fine-tuning. This allows applying fine-tuned adaptations to the base model.

```python
wrapper = CosmosPredictWrapper(
    extract_layer=19,
    lora_path='./cosmos-lora-adapter'
)
```

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

### Main Model Import

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/modules-architecture.md

Import the main MimicVideo class for use in your project.

```python
from mimic_video import MimicVideo
```

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

### Get Model Device

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Retrieves the device (CPU or GPU) where the model's parameters are located. This is useful for ensuring data is moved to the correct device before being processed by the model.

```python
# Get model device
device = next(model.parameters()).device
```

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

### Helper Function: get_discounted_returns

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/INDEX.txt

Example usage of the `get_discounted_returns` helper function, likely used within the FlowSteering class or for RL data processing. It calculates discounted returns from rewards.

```python
import torch

# Assuming get_discounted_returns is accessible, e.g., from mimic_video.flow_steering.utils
# from mimic_video.flow_steering.utils import get_discounted_returns

def get_discounted_returns(rewards, discount_factor):
    # Placeholder for the actual function implementation
    pass

# Dummy rewards and discount factor
rewards = torch.tensor([1.0, 0.5, 0.2])
discount_factor = 0.99

# Calculate discounted returns
# discounted_returns = get_discounted_returns(rewards, discount_factor)
# print(discounted_returns)

```

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

### Initialize FlowSteering for Offline RL

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Wrap a trained MimicVideo model with FlowSteering for offline Reinforcement Learning. Configure parameters like discount factor, EMA decay, and expectile tau.

```python
# Wrap trained model with RL
rl_model = FlowSteering(
    model,
    discount_factor=0.99,
    ema_decay=0.995,
    expectile_tau=0.25,  # Pessimistic for offline
    use_minto=True
)
```

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

### Mimic-Video Architecture Overview

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/README.md

Illustrates the data flow and processing stages within the Mimic-Video model, from input to output, including optional video encoding and action processing steps.

```text
Input
├── Video (batch, frames, 3, h, w) or pre-encoded features
├── Joint State (batch, dim)
└── Action (batch, steps, dim)
    ↓
Video Encoding (optional)
└── CosmosPredictWrapper
    └── → (batch, seq_len, dim_hidden)
    ↓
Action Processing
├── Embedding + RNN
├── 8 × Transformer Layers
│   ├── Cross-Attention (with video)
│   ├── Self-Attention
│   └── SwiGLU Feedforward
└── Final Norm + Prediction Head
    ↓
Output
├── Training: MSE Flow Loss
└── Inference: Predicted Flow
```

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

### MimicVideo Configuration: Large (Maximum Capacity)

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Initialize MimicVideo with large dimensions, increased depth, and a specific configuration for extracted video layer indices, aiming for maximum capacity.

```python
MimicVideo(dim=768, depth=12, action_chunk_len=64, extracted_video_layer_indices=[0,1,1,2,2,2,...])
```

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

### RL Model Usage: Multi-step Returns

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/flow-steering-api.md

Example of using the RL model for multi-step returns with variable trajectory lengths. Rewards should have shape (batch, n_steps) and n_step_lens should specify actual lengths.

```python
# Multi-step returns with variable lengths
loss, breakdown = rl_model(
    prompts=['task1', 'task2'],
    video=video,
    joint_state=joint_state,
    actions=actions,
    noise_latents=noise_latents,
    next_video=next_video,
    next_joint_state=next_joint_state,
    rewards=torch.randn(2, 4),  # shape: (batch, 4 steps)
    n_step_lens=torch.tensor([3, 4])  # Actual lengths
)

loss.backward()
optimizer.step()
```

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

### Helper Function: expectile_l2_loss

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/INDEX.txt

Example usage of the `expectile_l2_loss` helper function, likely used in training stability or specific loss calculations within FlowSteering. It computes an L2 loss based on expectiles.

```python
import torch

# Assuming expectile_l2_loss is accessible, e.g., from mimic_video.flow_steering.utils
# from mimic_video.flow_steering.utils import expectile_l2_loss

def expectile_l2_loss(pred, target, expectile=0.8):
    # Placeholder for the actual function implementation
    pass

# Dummy prediction and target tensors
pred = torch.randn(10)
target = torch.randn(10)

# Calculate expectile L2 loss
# loss = expectile_l2_loss(pred, target, expectile=0.8)
# print(loss)

```

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

### CosmosPredictWrapper.finetune

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/cosmos-predict-api.md

Fine-tunes the Cosmos transformer using LoRA on custom robot video data. This method handles the setup for LoRA adaptation, including optional full transformer unfreezing and distributed training acceleration.

```APIDOC
## finetune

### Description
Fine-tune the Cosmos transformer using LoRA (Low-Rank Adaptation) on custom robot video data. This method configures the model for fine-tuning, allowing for efficient adaptation to specific datasets.

### Method Signature
```python
def finetune(
    self,
    dataset: torch.utils.data.Dataset,
    save_path: str = "cosmos-lora-adapter",
    batch_size: int = 1,
    lr: float = 1e-4,
    epochs: int = 1,
    mu: float = 0.0,
    sigma: float = 1.0,
    lora_config: dict | LoraConfig | None = None,
    accelerator: Accelerator | None = None,
    unfreeze_transformer: bool = False,
    train_fixed_video_prefix_max_delay: int = 0
) -> None
```

### Parameters

#### Dataset
- **dataset** (Dataset) - Required - PyTorch Dataset yielding (video, text) tuples. video shape (num_frames, 3, height, width) in [0, 1] range, text is str or tokenized.

#### Training Configuration
- **save_path** (str) - Optional - Default: "cosmos-lora-adapter" - Directory to save trained LoRA adapter weights.
- **batch_size** (int) - Optional - Default: 1 - Batch size for training.
- **lr** (float) - Optional - Default: 1e-4 - AdamW learning rate.
- **epochs** (int) - Optional - Default: 1 - Number of training epochs.
- **mu** (float) - Optional - Default: 0.0 - Mean of logit-normal timestep sampling.
- **sigma** (float) - Optional - Default: 1.0 - Std of logit-normal timestep sampling.

#### LoRA and Accelerator
- **lora_config** (dict | LoraConfig | None) - Optional - Default: None - LoRA configuration. If dict, passed to peft.LoraConfig. If None, uses DEFAULT_LORA_CONFIG.
- **accelerator** (Accelerator | None) - Optional - Default: None - Huggingface Accelerator for distributed training. If None, creates one internally.

#### Advanced Options
- **unfreeze_transformer** (bool) - Optional - Default: False - If True, fully unfreeze transformer parameters (full fine-tuning). If False, only LoRA adapters are trainable.
- **train_fixed_video_prefix_max_delay** (int) - Optional - Default: 0 - Maximum prefix length for training-time prefixing. 0 = disabled.

### Behavior

1. Wraps transformer with LoRA adapters (or unfreezes if unfreeze_transformer=True).
2. Encodes videos with VAE (frozen), extracts text embeddings (frozen).
3. Samples timesteps and applies flow matching loss with optional video prefixing.
4. Saves trained LoRA weights to save_path.

```

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

### Distributed Training with PyTorch Lightning

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Configures a PyTorch Lightning `Trainer` for distributed training across multiple devices using the 'ddp' strategy. PyTorch Lightning simplifies distributed training setup and management.

```python
# Using PyTorch Lightning (recommended)
from pytorch_lightning import Trainer
trainer = Trainer(devices=4, strategy='ddp')
```

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

### MimicVideo Data Flow (Training)

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/modules-architecture.md

Illustrates the data flow during the training process for the MimicVideo model, from input tensors to loss calculation.

```python
Input Tensors:
├── actions (b, na, da)
├── joint_state (b, dj)
└── video/video_hiddens (b, n, dv)
    ↓
Embedding & RNN:
├── action_tokens = to_action_tokens(actions)
├── action_rnn_out = rnn(action_tokens)
└── joint_token = to_joint_state_token(joint_state)
    ↓
Time Conditioning:
└── time_cond = to_time_cond(fourier_embed(time))
    ↓
Transformer Stack:
├── For each layer:
│   ├── cross_attn(actions, context=video)
│   ├── self_attn(actions)
│   └── feedforward(actions)
│   └── [Each with gated residual aggregation]
    ↓
Final Norm & Prediction:
├── embed = final_norm(tokens)
└── pred = to_pred(embed)
    ↓
Loss:
├── If training: flow_loss = MSE(pred_flow, target_flow)
└── If inference: return flow for Euler stepping
```

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

### FlowSteering Sample for Policy Evaluation

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/INDEX.txt

Demonstrates the use of the sample method in FlowSteering for policy evaluation. This method typically takes states and returns sampled actions.

```python
import torch
from mimic_video.flow_steering.flow_steering import FlowSteering

model = FlowSteering(
    num_tokens = 8192,
    dim = 512,
    depth = 6,
    heads = 8
)

# Dummy state data
states = torch.randn(1, 10)

# Sample actions from the policy
# sampled_actions = model.sample(states) # Example call, actual call signature may vary

```

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

### Medium Model Configuration

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/configuration.md

Configures a medium-sized Mimic Video model for a balance between performance and speed. Uses a pretrained Cosmos model.

```python
wrapper = CosmosPredictWrapper(
    model_name='nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
    extract_layer=19
)

model = MimicVideo(
    dim=512,
    video_predict_wrapper=wrapper,
    depth=8,
    action_chunk_len=32,
    dim_action=20,
    dim_joint_state=32,
    train_time_rtc=True,
    train_time_rtc_max_delay=8
)
```

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

### MimicVideo Constructor

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/mimic-video-api.md

Initializes the MimicVideo model. This constructor accepts numerous parameters to configure the model's dimensions, attention mechanisms, time conditioning, and output processing for video-action prediction.

```python
def __init__(
    self,
    dim: int,
    video_predict_wrapper: Module | None = None,
    *,
    dim_video_hidden: int | None = None,
    action_chunk_len: int = 32,
    dim_action: int = 20,
    dim_joint_state: int = 32,
    proprio_mask_prob: float = 0.1,
    depth: int = 8,
    dim_head: int = 64,
    heads: int = 8,
    attention_residual_heads: int = 4,
    attention_residual_dim_head: int = 64,
    expansion_factor: float = 4.,
    is_flow_model: bool = True,
    ada_ln_zero_bias: float = -5.,
    dim_time_cond: int | None = None,
    pred_head: Module | None = None,
    sample_time_fn: Callable | None = None,
    train_time_rtc: bool = False,
    train_time_rtc_max_delay: int | None = None,
    action_mean_std: Tensor | None = None,
    joint_mean_std: Tensor | None = None,
    model_output_clean: bool = True,
    num_task_ids: int = 0,
    num_advantage_ids: int = 0,
    advantage_cfg_dropout: float = 0.25,
    extracted_video_layer_indices: list[int] | None = None,
    num_video_viewpoints: int = 1,
    video_time_denoise_mu: float = 0.,
    video_time_denoise_sigma: float = 1.,
    state_autoencoder: Module | dict | None = None,
    state_autoencoder_loss_weight: float = 1.,
    state_autoencoder_video_layer_index: int | None = None,
    eps: float = 1e-5
) -> None
```

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

### Initialize MimicVideo Model

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Initialize the MimicVideo model with specified dimensions and configurations. This constructor allows for optional video prediction wrappers and various training-time options.

```python
model = MimicVideo(
    dim=512,                              # Hidden dimension
    video_predict_wrapper=wrapper,        # Video encoder (optional)
    dim_video_hidden=4096,                # Video feature dim (if no wrapper)
    action_chunk_len=32,                  # Action sequence length
    dim_action=20,                        # Action dimensionality
    dim_joint_state=32,                   # Proprioceptive state dim
    depth=8,                              # Number of transformer layers
    train_time_rtc=True,                  # Real-time chunking
    train_time_rtc_max_delay=8,           # Max prefix for RTC
    action_mean_std=action_stats,         # Action normalization
    num_task_ids=0,                       # Number of task classes
    num_advantage_ids=0,                  # Number of return quantiles
    num_video_viewpoints=1,               # Multi-view camera count
    state_autoencoder={...},              # Auxiliary state encoder
)
```

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

### Task-Conditioned Policies

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/INDEX.txt

Shows how to set up mimic-video for task-conditioned policies, where the model's behavior is adapted based on a given task. This typically involves providing task embeddings or conditioning information.

```python
from mimic_video.mimic_video import MimicVideo

model = MimicVideo(
    num_tokens = 8192,
    dim = 512,
    depth = 6,
    heads = 8,
    # task_conditioning_enabled = True # Example parameter, actual parameter name may vary
)

# Dummy video and task data
video = torch.randn(1, 3, 128, 128)
task_embedding = torch.randn(1, 64)

# Forward pass with task conditioning
# loss = model(video, task_embedding) # Example call, actual call signature may vary

```

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

### CosmosPredictWrapper Initialization and Training Mode Usage

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/cosmos-predict-api.md

Initialize CosmosPredictWrapper for feature extraction. Demonstrates single-step feature extraction during training.

```python
import torch
from mimic_video.cosmos_predict import CosmosPredictWrapper

wrapper = CosmosPredictWrapper(extract_layer=19)

# Training: single-step feature extraction
video = torch.rand(2, 5, 3, 256, 256)  # 2 batches, 5 frames
video_hiddens = wrapper(
    videos=video,
    prompts=['rotate the object', 'grasp the object'],
    timestep=0.5  # Mid-denoising
)
# Output: (2, seq_len, 4096) - depends on VAE compression ratio
```

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

### Initialize MimicVideo for Offline RL

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Configures the MimicVideo model for offline Reinforcement Learning by enabling advantage conditioning and setting up a state autoencoder. A FlowSteering wrapper is then applied to the model.

```python
# Add advantage conditioning and state autoencoder
model = MimicVideo(
    dim=512,
    num_advantage_ids=5,
    state_autoencoder={'enc_depth': 2, 'dec_depth': 2, ...}
)
rl_model = FlowSteering(model, expectile_tau=0.25)
```

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

### Configuration Presets

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Provides pre-defined configurations for the MimicVideo model: Small, Medium, and Large.

```APIDOC
## Configuration Presets

Provides pre-defined configurations for the MimicVideo model: Small, Medium, and Large.

### Small (Testing)
```python
MimicVideo(dim=256, depth=4, action_chunk_len=16)
```

### Medium (Recommended)
```python
MimicVideo(dim=512, depth=8, action_chunk_len=32, train_time_rtc=True)
```

### Large (Maximum Capacity)
```python
MimicVideo(dim=768, depth=12, action_chunk_len=64, extracted_video_layer_indices=[0,1,1,2,2,2,...])
```
```

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

### MimicVideo Constructor

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Initializes the MimicVideo model with various configuration options for dimensionality, video encoding, action sequencing, and transformer layers.

```APIDOC
## MimicVideo Constructor

Initializes the MimicVideo model with various configuration options for dimensionality, video encoding, action sequencing, and transformer layers.

### Parameters

- **dim** (int) - Hidden dimension of the model.
- **video_predict_wrapper** (object, optional) - A wrapper for video encoding. Defaults to None.
- **dim_video_hidden** (int, optional) - Dimension of video features if no wrapper is provided.
- **action_chunk_len** (int) - The length of action sequences.
- **dim_action** (int) - The dimensionality of actions.
- **dim_joint_state** (int) - The dimensionality of the joint state.
- **depth** (int) - The number of transformer layers.
- **train_time_rtc** (bool) - Whether to use real-time chunking during training.
- **train_time_rtc_max_delay** (int, optional) - Maximum delay for real-time chunking.
- **action_mean_std** (tuple, optional) - Mean and standard deviation for action normalization.
- **num_task_ids** (int, optional) - Number of task classes.
- **num_advantage_ids** (int, optional) - Number of return quantiles.
- **num_video_viewpoints** (int, optional) - Number of multi-view camera inputs.
- **state_autoencoder** (dict, optional) - Configuration for an auxiliary state encoder.
```

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

### Initialize CosmosPredictWrapper with Random Weights

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/cosmos-predict-api.md

Instantiate the wrapper with random weights and a tiny architecture for testing or debugging purposes. This bypasses the need for pretrained model loading.

```python
wrapper = CosmosPredictWrapper(
    random_weights=True,
    tiny=True,
    extract_layer=1
)
```

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

### Using LoRA Config with CosmosPredictWrapper

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/types.md

Demonstrates how to apply LoRA fine-tuning to a model using the `CosmosPredictWrapper`. The `lora_config` is passed to the `finetune` method along with the dataset and number of epochs.

```python
wrapper = CosmosPredictWrapper(extract_layer=19)
wrapper.finetune(
    dataset=dataset,
    lora_config=lora_config,
    epochs=3
)
```

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

### Import Mimic-Video Components

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Import necessary classes from the mimic_video library for video prediction and related functionalities.

```python
from mimic_video import MimicVideo
from mimic_video.cosmos_predict import CosmosPredictWrapper
from mimic_video.flow_steering import FlowSteering
```

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

### Mimic Video Inference with Caching

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/types.md

Demonstrates how to perform inference with the MimicVideo model, utilizing the `Cache` object for incremental generation. The `Cache` is automatically managed by the `forward` method.

```python
import torch
from mimic_video.mimic_video import MimicVideo, Cache

model = MimicVideo(dim=512)

# During inference with caching
cache = None
for step in range(16):
    output, cache = model.forward(
        actions=current_actions,
        joint_state=joint_state,
        video_hiddens=video_features,
        time=torch.tensor(step / 16),
        cache=cache,
        return_cache=True
    )

# Cache is automatically managed by forward()
```

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

### Passing Custom Time Sampling Function to MimicVideo

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/types.md

Shows how to instantiate the `MimicVideo` model with a custom time sampling function, such as `default_sample_time_fn`.

```python
model = MimicVideo(
    dim=512,
    sample_time_fn=default_sample_time_fn  # or custom function
)
```

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

### Create RL Critic Wrapper

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Creates a critic component from the MimicVideo model, used for estimating state-action values. The critic is initialized with a specified depth and requires prompts, video, joint state, and actions as input to output a value estimate.

```python
# Critic: outputs value estimate
critic = model.create_critic_from(depth=6)
value = critic(prompts=..., video=..., joint_state=..., actions=...)
```

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

### Sample Action Sequences (Basic)

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/mimic-video-api.md

Basic sampling of action sequences. Requires specifying the number of steps, video, joint state, and prompts.

```python
# Basic sampling
actions = model.sample(
    steps=16,
    video=video,
    joint_state=joint_state,
    prompts='put the object down'
)
```

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

### Evaluation with Mimic Video

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/flow-steering-api.md

This snippet shows how to evaluate the Mimic Video model by sampling actions. It sets the model to evaluation mode and uses torch.no_grad() for efficiency. The `sample` method is used with specific prompts and observation data.

```python
# 4. Evaluation
rl_model.eval()
with torch.no_grad():
    test_actions, _ = rl_model.sample(
        prompts='task description',
        video=test_obs,
        joint_state=test_joint_state,
        exploration_noise_std=0.1
    )
```

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

### Passing State Autoencoder Config to MimicVideo

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/types.md

Shows how to instantiate the `MimicVideo` model with a custom `state_autoencoder` configuration dictionary.

```python
model = MimicVideo(
    dim=512,
    state_autoencoder=state_autoencoder
)
```

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

### Sample Action Sequences (Inpainting)

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/mimic-video-api.md

Inpainting action sequences by providing a prefix of actions to fix and regenerating the rest. Useful for continuing or correcting partial sequences.

```python
# Inpainting: fix first few actions, regenerate rest
prefix = torch.randn(1, 8, 20)  # Fix first 8 actions
actions = model.sample(
    steps=16,
    prefix_action_chunk=prefix,
    video=video,
    joint_state=joint_state,
    prompts='continue task'
)
```

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

### RL Training Loop

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/usage-patterns.md

Demonstrates a typical reinforcement learning training loop for the Mimic Video model. It includes computing discounted returns, calculating RL loss, performing backpropagation with gradient clipping, and updating Exponential Moving Average (EMA) weights.

```python
rl_model.train()
for episode in rl_dataset:
    # Compute discounted returns
    with torch.no_grad():
        next_actions, next_noise = rl_model.sample(
            prompts=episode.next_prompts,
            video=episode.next_video,
            joint_state=episode.next_joint_state
        )
        next_q = rl_model.outer_critic(
            prompts=episode.next_prompts,
            video=episode.next_video,
            joint_state=episode.next_joint_state,
            actions=next_actions
        )
    
    target_returns = get_discounted_returns(
        rewards=episode.rewards,
        last_q=next_q,
        discount_factor=0.99,
        done=episode.done
    )
    
    # RL loss
    loss, (actor_loss, q_loss, steer_loss) = rl_model(
        prompts=episode.prompts,
        video=episode.video,
        joint_state=episode.joint_state,
        actions=episode.actions,
        noise_latents=torch.randn_like(episode.actions),
        next_video=episode.next_video,
        next_joint_state=episode.next_joint_state,
        rewards=episode.rewards,
        done=episode.done
    )
    
    # Backward
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(
        list(rl_model.actor.parameters()) + 
        list(rl_model.outer_critic.parameters()) + 
        list(rl_model.inner_critic.parameters()),
        1.0
    )
    optimizer.step()
    
    # Update EMA
    rl_model.update_critic_ema()
    
    print(f"A:{actor_loss:.3f} Q:{q_loss:.3f} S:{steer_loss:.3f}")

# Step 4: Evaluation with exploration
rl_model.eval()
actions, noise = rl_model.sample(
    prompts='task',
    video=test_video,
    joint_state=test_joint_state,
    exploration_noise_std=0.1
)
```

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

### Prepare Data for Mimic Video Training

Source: https://github.com/lucidrains/mimic-video/blob/main/README.md

Prepares sample video, joint state, and action tensors for training the MimicVideo model. The video tensor shape is (batch_size, frames, channels, height, width).

```python
# states

video = torch.rand(2, 5, 3, 32, 32) # 5 frames, 3 channels, 32 x 32

joint_state = torch.randn(2, 32)

# action

actions = torch.randn(2, 32, 20)
```

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

### Inference Optimization with Caching

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/INDEX.txt

Shows how to optimize inference by enabling caching mechanisms within MimicVideo. Caching can significantly speed up predictions by reusing previously computed results.

```python
from mimic_video.mimic_video import MimicVideo

model = MimicVideo(
    num_tokens = 8192,
    dim = 512,
    depth = 6,
    heads = 8,
    # caching_enabled = True # Example parameter, actual parameter name may vary
)

# Dummy data for inference
video = torch.randn(1, 3, 128, 128)

# Perform inference with caching enabled
# with torch.no_grad():
#     output = model.sample(video) # Example call, actual call signature may vary

```

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

### Create RL Actor Wrapper

Source: https://github.com/lucidrains/mimic-video/blob/main/_autodocs/QUICK-REFERENCE.md

Creates an actor component from the MimicVideo model, which is responsible for outputting policy distributions. The actor is initialized with a specified depth and takes prompts, video, and joint state as input to produce mean and log variance.

```python
# Actor: outputs policy distribution
actor = model.create_actor_from(depth=6)
mean, logvar = actor(prompts=..., video=..., joint_state=...)
```