RoboImagine

Learning Robotic World Models with Semantic and Visual Feedback

This repository provides inference checkpoints, FSQ visual codecs, action normalization ranges and six small reproducibility examples. The GitHub repository contains the model implementation, memory/search code, inference notebook and experimental figure/table reproduction.

Released checkpoints

Dataset	main checkpoint	sft checkpoint	Action values per step
RT-1	GDBS + semantic/visual feedback, 500 updates	Matched SFT initialization	13
BridgeV2 (bridge)	Validation-selected checkpoint, 250 updates	Matched text-enabled SFT initialization	13
RoboCasa	GDBS + semantic/visual feedback, 500 updates	Matched SFT initialization	52

Each model includes its required trained text adapters. Each dataset includes its visual codec and action ranges. The frozen T5-base encoder is downloaded separately from google-t5/t5-base at the revision specified in the configuration. Optimizer states are not included.

The files use flat names such as rt1--main--model.safetensors. The public loader assembles the directory structure automatically and verifies every file's SHA-256 against release_manifest.json. Tensor files are copied without quantization or modification from the corresponding experiments.

Quick start

git clone https://github.com/Alexander-wu/RoboImagine.git
cd RoboImagine
# Install a compatible CUDA PyTorch/torchvision pair first.
python -m pip install -e '.[notebook]'
roboimagine --dataset rt1 --variant main --example-index 0 --output predictions/rt1

Or open Inference.ipynb and Run All. Use --dataset bridge or --dataset robocasa, and --variant sft to compare the matching initialization. Detailed input specifications, version information and output interpretation are in the inference guide.

Prediction takes two initial RGB observations, the planned action sequence and an instruction. Future ground-truth frames are used only for optional evaluation after prediction. The default decoder uses width-4 beam search and a fixed initial context, fixed dynamic anchor and six recent states. Longer videos require longer action sequences; the six-state window does not limit output length to six frames.

Examples and interpretation

For each dataset, example000 and example001 are the first two entries of the archived reporting/evaluation list. These are demonstration inputs, not a statistically representative benchmark. Each NPZ contains two uint8 RGB observations at 256×320, 33 float32 action entries, the instruction and 32 ground-truth frames. expected000/001 are separate saved float-pixel predictions and tokens from the main checkpoint, used to verify implementation parity.

Bridge's checkpoint is chosen on a 32-trajectory validation cohort. The manuscript's corresponding reported Bridge row uses that validation cohort; the released demonstration examples come from its separate reporting pool. Do not infer a full-test-set average from the validation score or the two demonstrations. Other experimental cohorts, including the paired 30-trajectory horizon/window study, are documented independently in the code repository.

These are offline predictive world models. The release does not establish closed-loop robot-control safety or task success. A complete training launcher, reward-service hosting and full raw training datasets are outside this inference release.

Attribution and provenance

• RT-1 / Fractal: Brohan et al., Robotics at Google. Official project, TFDS catalogue.
• BridgeData V2: Walke et al., CoRL 2023. Official dataset, CC BY 4.0.
• RoboCasa / RoboCasa365: Nasiriany et al., RSS 2024 / ICLR 2026. Official repository, datasets/assets CC BY 4.0.

The examples preserve the archived resizing/temporal/action preprocessing; original dataset authors retain their rights. Source IDs and hashes are in the manifest. The implementation uses Transformers/LLaMA, frozen T5-base, and an iVideoGPT/diffusers/FSQ-derived visual codec. Upstream notices and license texts are retained on GitHub. No blanket relicensing of third-party code, datasets or assets is asserted. See data and provenance.