RISE-pytorch 0.0.1

RISE (wip)

Implementation of RISE, Self-Improving Robot Policy with Compositional World Model

Usage

Here is a complete, minimal example of how to orchestrate self-improvement using RISE.

This involves:

1. Fine-tuning an action-conditioned Cosmos Dynamics Model.
2. Initializing a PiZero policy and SigLIP Value Network.
3. Engaging the RISE loop: Imagination Rollout followed by Policy Finetuning.

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

Citations

@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}, 
}