FASTER for Robomimic Images and Action Chunks

This repository extends the official FASTER: Value-Guided Sampling for Fast RL Robomimic implementation with image + proprioception observations and receding-horizon action chunks. The maintained and evaluated path uses FasterEXPOLearner.

Code References

The three revisions below were checked against the current implementation. FASTER-VLA and EXPO-FT are implementation references.

Project	Revision	Relationship
FASTER	7214f3b	Exact upstream base
FASTER-VLA	9cff3db	External check for image-conditioned VLA inputs and action-horizon/replanning interfaces. This repository does not use its π0.5/OpenPI model, but uses a from-scratch Robomimic policy
EXPO-FT	7a39b3c	Implementation reference for executed-prefix editing, C-step targets, validity weighting, and preservation of the unexecuted action tail

Changes from official FASTER

The official FASTER Robomimic implementation uses low-dimensional observations and single-step actions. This repository retains the H=C=1 path and adds:

Area	Change	Main files
Observations	Two RGB views plus 9D robot proprioception; Robomimic object state is not consumed in image mode	faster/networks/vision.py, faster/data/robomimic_action_chunk.py
Action chunks	Predict H actions, execute the first C, and keep chunks inside episode boundaries	faster/data/action_chunk.py, faster/data/replay_buffer.py, train_robo.py
EXPO-style updates	C-step returns, γ^C bootstrap, validity masks, and executed-prefix residual editing	faster/agents/faster_expo_learner.py
Image policy	Per-view ResNet-18 encoders and a temporal U-Net diffusion head	faster/networks/vision.py, faster/networks/conditional_unet1d.py
Data	Automatic raw demonstration download and image + proprioception HDF5 conversion	scripts/download_robomimic_datasets.py, scripts/run_robomimic_converter.py
Checkpoints	Atomic offline/online saves and validated warm-start loading from absolute, relative, or symlinked paths	train_robo.py

Image observations have this training interface:

{
    "state": float32[..., 9],
    "image": uint8[..., height, width, 6],
}

The two RGB views are channel-stacked at the data boundary and encoded separately. The actor and critic use separate, from-scratch visual encoders.

Validation results

These local FASTER-EXPO results use image + proprioception, H=16, C=8, seed 42, and 25 evaluation episodes. Robomimic's sparse evaluation return is reported as success rate. Each run uses 200k offline updates followed by 150k primitive online environment steps.

Task	Dataset	Best offline	Offline at 200k	Best online	Online at 150k
Lift	PH, 10 demos	96% at 50k	76%	92% at 50k	88%
Can	MH, all demos	100% at 50k/100k	92%	92% at 50k	88%
Square	PH, all demos	88% at 100k	76%	84% at 50k	76%

These are single-seed results with 4-percentage-point evaluation resolution. They verify that the end-to-end path works, but they are not statistically robust benchmarks.

A known numerical issue remains: the validation runs reached temperature ≈ 4.6e17 and edit_actor_loss ≈ 3.1e19. Only FasterEXPOLearner was evaluated; no IDQL result is reported.

The W&B curves show that temperature and edit-actor loss grow sharply around 100k offline updates and remain near those magnitudes during online training. Entropy stays near -68. All plots compare the same Lift, Can, and Square runs used in the table above.

Metric	Offline training	Online training
Temperature
Edit actor loss
Entropy

Raw W&B exports are stored beside the figures in docs/assets/.

Setup

Follow the official FASTER setup. This fork follows the official uv workflow. Its checked-in lockfile pins the Robomimic and Robosuite revisions and uses cmake<4 for compatibility with the legacy egl-probe build. From the repository root, only the fork-specific setup is:

uv sync --locked

export ROBOMIMIC_DATASETS_PATH=/absolute/path/to/robomimic
uv run --locked wandb login

The canonical launchers enable online W&B logging. Large datasets, uv caches,virtual environments, and exp/ may be placed outside the checkout using ROBOMIMIC_DATASETS_PATH, UV_CACHE_DIR, UV_PROJECT_ENVIRONMENT, and an optional exp symlink.

Dataset preparation

The documented image launchers expect:

${ROBOMIMIC_DATASETS_PATH}/
├── lift/ph/image_v141.hdf5
├── can/mh/image_v141.hdf5
└── square/ph/image_v141.hdf5

All three tasks use agentview and robot0_eye_in_hand images stored at 84 × 84 and cropped to 76 × 76 during training. The dataset also contains proprioception; the image training loader selects the 9-D robot state and does not use the object observation.

Each launcher checks its dataset before training. When image_v141.hdf5 is missing, it downloads the registered raw demonstration file and runs Robomimic's dataset_states_to_obs.py converter. Existing valid datasets are left untouched. Image generation requires a working MuJoCo EGL device.

Manual preparation:

# Lift and Square use PH.
uv run --locked python scripts/download_robomimic_datasets.py \
  --envs lift square \
  --splits ph \
  --hdf5-types image \
  --compress-images

# The canonical Can launcher uses MH.
uv run --locked python scripts/download_robomimic_datasets.py \
  --envs can \
  --splits mh \
  --hdf5-types image \
  --compress-images

Set FASTER_AUTO_PREPARE_DATASETS=0 to disable automatic preparation and fail when data is missing. Image training defaults to --image_cache_mode=all; use --image_cache_mode=low_dim when host RAM is insufficient.

Training

Task	Launcher	Dataset
Lift	bash scripts/faster_expo_image_lift.sh	PH, 10 demos
Can	bash scripts/faster_expo_image_can.sh	MH, all demos
Square	bash scripts/faster_expo_image_square.sh	PH, all demos

Common launcher settings:

Setting	Value
Prediction / execution horizon	H=16 / C=8
Action head	temporal U-Net, base width 128
Diffusion / candidates	T=10 / N=8
Batch size	64
Offline bootstrap	200k updates, UTD 1
Online training	150k primitive steps, UTD 20
Offline ratio	0.5
Replay memory budget	4 GiB
Evaluation	25 episodes every 50k steps
W&B mode	online

Checkpoints

Launchers save phase-separated checkpoints:

exp/<run>/checkpoints/
├── offline/agent_<update>.msgpack
└── online/agent_<environment-step>.msgpack

Warm-start a compatible checkpoint:

bash scripts/faster_expo_image_square.sh \
  --restore_path=/absolute/path/to/agent_200000.msgpack

Relative restore paths are resolved from the repository root. Restore skips offline pretraining by default; pass --skip_pretrain_on_restore=False to run it again. Replay contents, environment state, and W&B state are not restored.

Papers and References

• FASTER paper
• FASTER official code
• FASTER-VLA official code
• EXPO-FT paper