Complete checkpoint release with codecs, examples, manifests and verification

cb4f5e2 verified 19 days ago

4.82 kB

---
tags:
- robotics
- world-model
- video-prediction
- autoregressive
- safetensors
language:
- en
---
# 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](https://github.com/Alexander-wu/RoboImagine) 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

```bash
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](https://github.com/Alexander-wu/RoboImagine/blob/main/notebooks/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](https://github.com/Alexander-wu/RoboImagine/blob/main/docs/INFERENCE.md).

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](https://robotics-transformer1.github.io/), [TFDS catalogue](https://www.tensorflow.org/datasets/catalog/fractal20220817_data).
- BridgeData V2: Walke et al., CoRL 2023. [Official dataset](https://rail-berkeley.github.io/bridgedata/), CC BY 4.0.
- RoboCasa / RoboCasa365: Nasiriany et al., RSS 2024 / ICLR 2026. [Official repository](https://github.com/robocasa/robocasa), 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](https://github.com/Alexander-wu/RoboImagine/blob/main/docs/DATA_AND_PROVENANCE.md).

	---
	tags:
	- robotics
	- world-model
	- video-prediction
	- autoregressive
	- safetensors
	language:
	- en
	---

	# 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](https://github.com/Alexander-wu/RoboImagine) 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

	```bash
	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](https://github.com/Alexander-wu/RoboImagine/blob/main/notebooks/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](https://github.com/Alexander-wu/RoboImagine/blob/main/docs/INFERENCE.md).

	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](https://robotics-transformer1.github.io/), [TFDS catalogue](https://www.tensorflow.org/datasets/catalog/fractal20220817_data).
	- BridgeData V2: Walke et al., CoRL 2023. [Official dataset](https://rail-berkeley.github.io/bridgedata/), CC BY 4.0.
	- RoboCasa / RoboCasa365: Nasiriany et al., RSS 2024 / ICLR 2026. [Official repository](https://github.com/robocasa/robocasa), 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](https://github.com/Alexander-wu/RoboImagine/blob/main/docs/DATA_AND_PROVENANCE.md).