RISE

https://github.com/lucidrains/rise
Admin
RISE is a self-improving robot policy framework that uses a compositional world model for...
Tokens:5,779
Snippets:37
Trust Score:10
License:MIT
Update:3 months ago
Show doc for...
Tokens:
Raw
### Train Dynamics Model with DynamicsTrainer

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

Train the dynamics model using the DynamicsTrainer. This includes starting the training loop or performing a single training step.

```python
# Train the dynamics model
trainer.train(num_steps = 1000)

# Manual training step
batch = next(iter(trainer.dataloader))
loss = trainer.train_step(
    videos = batch['video'],
    next_videos = batch['next_video'],
    actions = batch['actions'],
    prompt_token_ids = batch['prompt_token_ids']
)
```

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

### Initialize DynamicsTrainer

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

Initialize the DynamicsTrainer with a dynamics model, dataset, and training parameters. Supports Accelerate for distributed training.

```python
from RISE import CosmosPredictWrapper, DynamicsTrainer, MockOfflineRoboticFrameDataset

# Create dynamics model
dynamics_model = CosmosPredictWrapper(
    model_name = 'nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
    action_dim = 10,
    random_weights = True,
    tiny = True
)

# Create offline dataset
dataset = MockOfflineRoboticFrameDataset(
    num_samples = 100,
    prefix_frames = 8,
    image_size = 64,
    action_dim = 10
)

# Initialize trainer with Accelerate support
trainer = DynamicsTrainer(
    model = dynamics_model,
    dataset = dataset,
    batch_size = 2,
    lr = 1e-4,
    weight_decay = 0.0,
    max_grad_norm = 0.5,
    shuffle = True,
    cpu = False,                      # Set True for CPU-only training
    accelerate_kwargs = {}            # Additional Accelerate config
)
```

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

### Initialize CosmosPredictWrapper

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

Initialize the CosmosPredictWrapper with pretrained weights or random weights for testing. Specify action dimension, layers to extract, and data type.

```python
from RISE import CosmosPredictWrapper

# Initialize with pretrained Cosmos model
dynamics_model = CosmosPredictWrapper(
    model_name = 'nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
    action_dim = 10,                  # Dimension of action space
    extract_layers = 19,              # Transformer layers to extract features from
    torch_dtype = torch.bfloat16      # Model precision
)

# For testing/development, use random weights with tiny architecture
dynamics_model_tiny = CosmosPredictWrapper(
    model_name = 'nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
    action_dim = 10,
    random_weights = True,            # Use random initialization
    tiny = True                       # Use small model for testing
)
```

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

### Orchestrate RISE Self-Improvement Loop

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

Demonstrates the complete workflow including dynamics model training, policy initialization, imagination rollouts, and advantage-conditioned fine-tuning.

```python
import torch

from RISE import (
    RISE,
    CosmosPredictWrapper,
    DynamicsTrainer,
    MockOfflineRoboticFrameDataset
)

from value_network import SigLIPValueNetwork

from pi_zero_pytorch.pi_zero import PiZero, SigLIP as PiZeroSigLIP

# 1. Provide an offline seed dataset (video + proprioception)
dataset = MockOfflineRoboticFrameDataset(num_samples = 10, image_size = 224)

# 2. Initialize and Fine-Tune the Action-Conditioned Dynamics Model
dynamics_model = CosmosPredictWrapper(
    model_name = 'nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
    action_dim = 10,
    action_chunk_len = 8
)

trainer = DynamicsTrainer(
    model = dynamics_model,
    dataset = dataset,
    batch_size = 1,
    lr = 1e-4
)

# Learn system dynamics from the offline dataset
trainer.train(num_steps = 1000)

# 3. Initialize Policy (PiZero) and Value Network Evaluator
policy = PiZero(
    dim = 256,
    num_tokens = 1000,
    dim_action_input = 10,
    dim_joint_state = 14,
    depth = 4,
    pi05 = True,
    vit = PiZeroSigLIP(
        image_size = 224,
        patch_size = 16,
        dim = 256,
        depth = 4,
        heads = 4,
        mlp_dim = 512
    ),
    vit_dim = 256
)

value_model = SigLIPValueNetwork(
    siglip_image_size = 224,
    siglip_patch_size = 16,
    siglip_dim = 256,
    siglip_depth = 4,
    siglip_heads = 4,
    siglip_mlp_dim = 512
)

# 4. Instantiate the RISE Orchestrator
rise = RISE(
    policy = policy,
    dynamics_model = dynamics_model,
    value_model = value_model,
    trajectory_length = 8,
    num_prompt_tokens = 12,
    imagination_steps = 5
)

# 5. Imagination Rollout Stage
# The policy generates actions, the dynamics model predicts the future,
# the value network evaluates the advantage, and experience is stored.
replay_buffer = rise.imagination_rollout(
    seed_dataset = dataset,
    num_episodes = 2,
    batch_size = 1,
    buffer_folder = './rise_experience_buffer'
)

# 6. Self-Improvement Finetuning Stage
# The policy learns to imitate high-advantage imagined trajectories.
rise.finetune_with_advantage_conditioning(
    replay_buffer = replay_buffer,
    num_steps = 1000,
    batch_size = 2,
    lr = 1e-4
)

# 7. Save the Self-Improved Policy
torch.save(rise.policy.state_dict(), './improved_pi05_policy.pt')
```

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

### Initialize RISE Orchestrator

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

Configures the policy network, dynamics model, and value network before instantiating the main RISE orchestrator.

```python
import torch
from RISE import RISE, CosmosPredictWrapper, MockOfflineRoboticFrameDataset
from value_network import SigLIPValueNetwork
from pi_zero_pytorch.pi_zero import PiZero, SigLIP as PiZeroSigLIP

# Initialize policy network (PiZero with SigLIP vision encoder)
policy = PiZero(
    dim = 256,
    num_tokens = 1000,
    dim_action_input = 10,
    dim_joint_state = 14,
    depth = 4,
    pi05 = True,
    vit = PiZeroSigLIP(
        image_size = 224,
        patch_size = 16,
        dim = 256,
        depth = 4,
        heads = 4,
        mlp_dim = 512
    ),
    vit_dim = 256
)

# Initialize dynamics model (action-conditioned Cosmos)
dynamics_model = CosmosPredictWrapper(
    model_name = 'nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
    action_dim = 10,
    action_chunk_len = 8
)

# Initialize value network
value_model = SigLIPValueNetwork(
    siglip_image_size = 224,
    siglip_patch_size = 16,
    siglip_dim = 256,
    siglip_depth = 4,
    siglip_heads = 4,
    siglip_mlp_dim = 512
)

# Create RISE orchestrator
rise = RISE(
    policy = policy,
    dynamics_model = dynamics_model,
    value_model = value_model,
    trajectory_length = 8,          # Length of action trajectories
    num_prompt_tokens = 12,         # Number of prompt tokens for conditioning
    imagination_steps = 5,          # Diffusion steps for action generation
    num_advantage_bins = 10,        # Bins for discretizing advantage values
    frames_per_future_latent = 5    # Frames per predicted latent
)
```

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

### Complete RISE Self-Improvement Pipeline

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

Demonstrates the full RISE loop: dataset creation, dynamics model training, policy/value network initialization, imagination rollouts, policy finetuning, and saving the improved policy.

```python
import torch
from RISE import (
    RISE,
    CosmosPredictWrapper,
    DynamicsTrainer,
    MockOfflineRoboticFrameDataset
)
from value_network import SigLIPValueNetwork
from pi_zero_pytorch.pi_zero import PiZero, SigLIP as PiZeroSigLIP

# Step 1: Create offline seed dataset
dataset = MockOfflineRoboticFrameDataset(
    num_samples = 10,
    image_size = 224
)

# Step 2: Initialize and fine-tune dynamics model
dynamics_model = CosmosPredictWrapper(
    model_name = 'nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
    action_dim = 10,
    action_chunk_len = 8
)

trainer = DynamicsTrainer(
    model = dynamics_model,
    dataset = dataset,
    batch_size = 1,
    lr = 1e-4
)
trainer.train(num_steps = 1000)

# Step 3: Initialize policy and value networks
policy = PiZero(
    dim = 256,
    num_tokens = 1000,
    dim_action_input = 10,
    dim_joint_state = 14,
    depth = 4,
    pi05 = True,
    vit = PiZeroSigLIP(
        image_size = 224, patch_size = 16,
        dim = 256, depth = 4, heads = 4, mlp_dim = 512
    ),
    vit_dim = 256
)

value_model = SigLIPValueNetwork(
    siglip_image_size = 224, siglip_patch_size = 16,
    siglip_dim = 256, siglip_depth = 4,
    siglip_heads = 4, siglip_mlp_dim = 512
)

# Step 4: Create RISE orchestrator
rise = RISE(
    policy = policy,
    dynamics_model = dynamics_model,
    value_model = value_model,
    trajectory_length = 8,
    num_prompt_tokens = 12,
    imagination_steps = 5
)

# Step 5: Imagination rollout
replay_buffer = rise.imagination_rollout(
    seed_dataset = dataset,
    num_episodes = 2,
    batch_size = 1,
    buffer_folder = './rise_experience_buffer'
)

# Step 6: Policy finetuning
rise.finetune_with_advantage_conditioning(
    replay_buffer = replay_buffer,
    num_steps = 1000,
    batch_size = 2,
    lr = 1e-4
)

# Step 7: Save improved policy
torch.save(rise.policy.state_dict(), './improved_pi05_policy.pt')
```

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

### Initialize MockOfflineRoboticFrameDataset

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

Initialize the MockOfflineRoboticFrameDataset for testing purposes. It generates mock data for training dynamics models.

```python
from RISE import MockOfflineRoboticFrameDataset
from torch.utils.data import DataLoader
```

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

### Perform Imagination Rollout

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

Collects imagined experiences by running the policy and dynamics model, storing results in a replay buffer.

```python
from RISE import RISE, CosmosPredictWrapper, MockOfflineRoboticFrameDataset

# Create a seed dataset with offline robot data
dataset = MockOfflineRoboticFrameDataset(
    num_samples = 100,
    prefix_frames = 8,
    image_size = 224,
    proprio_dim = 14,
    action_dim = 10
)

# Perform imagination rollout
# Returns a ReplayBuffer containing imagined experiences
replay_buffer = rise.imagination_rollout(
    seed_dataset = dataset,           # Offline dataset to seed imagination
    num_episodes = 10,                # Number of imagination episodes to collect
    batch_size = 2,                   # Batch size for rollout
    max_timesteps = 1,                # Autoregressive rollout steps per episode
    buffer_folder = './experience',   # Folder to store replay buffer
    predict_num_future_latents = 1,   # Number of future latents to predict
    force_new_buffer = True           # Clear existing buffer if present
)
```

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

### Create and Access Mock Robotic Dataset

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

Generates a mock dataset for robotic frames with specified dimensions and accesses individual samples. Useful for testing and development before using real data.

```python
dataset = MockOfflineRoboticFrameDataset(
    num_samples = 100,                # Total samples in dataset
    prefix_frames = 8,                # Video frames per sample
    channels = 3,                     # RGB channels
    image_size = 224,                 # Height and width
    proprio_dim = 14,                 # Proprioception dimensions
    action_dim = 10,                  # Action space dimensions
    action_chunk_len = 8,             # Action trajectory length
    num_prompt_tokens = 12,           # Task prompt tokens
    dtype = torch.bfloat16            # Data precision
)

# Access single sample
sample = dataset[0]
print(f"Video shape: {sample['video'].shape}")           # [8, 3, 224, 224]
print(f"Next video shape: {sample['next_video'].shape}") # [1, 3, 224, 224]
print(f"Actions shape: {sample['actions'].shape}")       # [9, 10]
print(f"Proprioception shape: {sample['proprioception'].shape}")  # [8, 14]
print(f"Prompt tokens shape: {sample['prompt_token_ids'].shape}") # [12]

# Use with DataLoader
loader = DataLoader(dataset, batch_size = 4, shuffle = True)
for batch in loader:
    videos = batch['video']          # [4, 8, 3, 224, 224]
    actions = batch['actions']       # [4, 9, 10]
    break
```

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

### Finetune Policy with Advantage Conditioning

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

Trains the policy network using the collected imagined trajectories and advantage values.

```python
# After imagination rollout, finetune the policy
rise.finetune_with_advantage_conditioning(
    replay_buffer = replay_buffer,    # Replay buffer from imagination_rollout
    num_steps = 1000,                 # Number of training steps
    batch_size = 4,                   # Training batch size
    lr = 1e-4,                        # Learning rate
    weight_decay = 0.0,               # Weight decay for optimizer
    max_grad_norm = 0.5               # Gradient clipping threshold
)
```

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

### CosmosPredictWrapper Flow Matching Loss

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

Compute the flow matching loss for fine-tuning the dynamics model. Requires current and next video frames, actions, and prompt tokens.

```python
from RISE import CosmosPredictWrapper

dynamics_model = CosmosPredictWrapper(
    model_name = 'nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
    action_dim = 10,
    random_weights = True,
    tiny = True
)

# Prepare training data
videos = torch.randn(2, 8, 3, 64, 64)       # Current video frames
next_videos = torch.randn(2, 1, 3, 64, 64)  # Target future frames
actions = torch.randn(2, 9, 10)              # Actions spanning both
prompt_token_ids = torch.randint(0, 1000, (2, 12))

# Compute flow matching loss for training
loss = dynamics_model(
    videos = videos,
    next_videos = next_videos,        # Target frames for supervision
    actions = actions,
    prompt_token_ids = prompt_token_ids,
    return_flow_loss = True           # Return loss instead of predictions
)

print(f"Flow matching loss: {loss.item():.4f}")
# Use loss.backward() for gradient updates
```

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

### Discretize Advantage Values with AdvantageBinner

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

Clamps continuous advantage values to a specified range and discretizes them into a fixed number of bins. Useful for advantage-conditioned policy learning.

```python
from RISE.RISE import AdvantageBinner
import torch

# Create advantage binner
binner = AdvantageBinner(
    min_val = -2.0,                   # Minimum advantage value
    max_val = 2.0,                    # Maximum advantage value
    num_bins = 10                     # Number of discrete bins
)

# Discretize advantage values
advantages = torch.tensor([0.5, -1.0, 1.5, -2.5, 3.0])
binned = binner(advantages)
print(f"Original advantages: {advantages}")
print(f"Binned indices: {binned}")
# Output: tensor([6, 2, 8, 0, 9])
# Values outside [-2, 2] are clamped to edge bins

# Use in policy conditioning
# Higher bin index = higher advantage = better trajectory
best_advantage_bin = binner.num_bins - 1  # = 9
```

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

### RISE Orchestrator Initialization

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

Initializes the RISE orchestrator by coordinating the policy network, dynamics model, and value network.

```APIDOC
## RISE Orchestrator Initialization

### Description
Initializes the main RISE class which manages the self-improvement loop, including the policy, dynamics model, and value network.

### Parameters
- **policy** (PiZero) - Required - The policy network for action generation.
- **dynamics_model** (CosmosPredictWrapper) - Required - The dynamics model for predicting future states.
- **value_model** (SigLIPValueNetwork) - Required - The value network for evaluating trajectories.
- **trajectory_length** (int) - Required - Length of action trajectories.
- **num_prompt_tokens** (int) - Required - Number of prompt tokens for conditioning.
- **imagination_steps** (int) - Required - Diffusion steps for action generation.
- **num_advantage_bins** (int) - Required - Bins for discretizing advantage values.
- **frames_per_future_latent** (int) - Required - Frames per predicted latent.
```

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

### RISE.finetune_with_advantage_conditioning

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

Trains the policy network using advantage-conditioned imitation learning on imagined trajectories.

```APIDOC
## RISE.finetune_with_advantage_conditioning

### Description
Performs advantage-conditioned imitation learning on the policy network using data from the replay buffer.

### Parameters
- **replay_buffer** (ReplayBuffer) - Required - Replay buffer from imagination_rollout.
- **num_steps** (int) - Required - Number of training steps.
- **batch_size** (int) - Required - Training batch size.
- **lr** (float) - Required - Learning rate.
- **weight_decay** (float) - Required - Weight decay for optimizer.
- **max_grad_norm** (float) - Required - Gradient clipping threshold.
```

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

### CosmosPredictWrapper Forward Pass

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

Perform a forward pass with the dynamics model to predict future video frames. Requires video, action, and prompt token inputs.

```python
# Forward pass for video prediction
videos = torch.randn(2, 8, 3, 224, 224)  # [batch, frames, channels, H, W]
actions = torch.randn(2, 8, 10)           # [batch, frames, action_dim]
prompt_token_ids = torch.randint(0, 1000, (2, 12))

# Predict future frames
predicted_video = dynamics_model(
    videos = videos,
    prompt_token_ids = prompt_token_ids,
    actions = actions,
    predict_num_future_latents = 1,   # Predict 1 future latent
    inference_steps = 15,             # Diffusion sampling steps
    return_decoded_latents = True     # Return decoded video frames
)
# Output shape: [batch, frames + future_frames, channels, H, W]
```

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

### RISE.imagination_rollout

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

Performs the imagination rollout stage where the policy generates actions and the dynamics model predicts future states.

```APIDOC
## RISE.imagination_rollout

### Description
Executes the imagination rollout process, collecting imagined experiences into a replay buffer for training.

### Parameters
- **seed_dataset** (MockOfflineRoboticFrameDataset) - Required - Offline dataset to seed imagination.
- **num_episodes** (int) - Required - Number of imagination episodes to collect.
- **batch_size** (int) - Required - Batch size for rollout.
- **max_timesteps** (int) - Required - Autoregressive rollout steps per episode.
- **buffer_folder** (str) - Required - Folder to store replay buffer.
- **predict_num_future_latents** (int) - Required - Number of future latents to predict.
- **force_new_buffer** (bool) - Required - Clear existing buffer if present.

### Response
- **replay_buffer** (ReplayBuffer) - Contains images, joint_state, actions, advantage_ids, and imagined_video.
```

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

### RISE BibTeX Citation

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

BibTeX entry for citing the RISE research paper.

```bibtex
@misc{yang2026riseselfimprovingrobotpolicy,
    title   = {RISE: Self-Improving Robot Policy with Compositional World Model}, 
    author  = {Jiazhi Yang and Kunyang Lin and Jinwei Li and Wencong Zhang and Tianwei Lin and Longyan Wu and Zhizhong Su and Hao Zhao and Ya-Qin Zhang and Li Chen and Ping Luo and Xiangyu Yue and Hongyang Li},
    year    = {2026},
    eprint  = {2602.11075},
    archivePrefix = {arXiv},
    primaryClass = {cs.RO},
    url     = {https://arxiv.org/abs/2602.11075}, 
}
```
Tokens:
Raw