Skip to content

Shawn-deeplearning/LARC

Repository files navigation

LARC

LARC stands for Latent Action Chunking with World-Model Rollout Consistency.

LARC trains fast goal-conditioned execution policies in the latent space of a frozen action-conditioned world model. It caches LeWorldModel latents once, trains either a single-step inverse dynamics policy (LARC-IDM) or a short action chunk policy (LARC-Chunk), and optionally fine-tunes the chunk policy with a frozen world-model rollout consistency loss.

The old warp Python package is kept as a compatibility layer. New code should import and run larc.

PushT Results

Evaluation setting: single seed, 50 PushT evaluation episodes, frozen LeWorldModel latents, execute_chunk_size=1, goal offsets of 25/35/50 environment steps. Evaluation time includes environment rollout and video saving overhead. Latency is the mean direct policy inference latency measured in the saved runs.

Method Goal offset 25 Goal offset 35 Goal offset 50 Mean latency
LARC-IDM 88% 82% 54% 14.9 ms
LARC-Chunk 98% 88% 70% 16.0 ms
LARC-Chunk + WM-FT 98% 90% 80% 15.9 ms

The main trend is that action chunking improves over single-step inverse dynamics, and frozen world-model rollout consistency improves long-horizon robustness, especially at goal offset 50.

Repository Layout

larc/
  world_model.py          # frozen LeWorldModel wrapper and WorldModel interface
  policies/               # IDM and action-chunk policy backbones
  losses.py               # BC, smoothness, and world-model consistency losses
  train_coral.py          # LARC policy training entrypoint
  eval_coral_pusht.py     # PushT evaluation entrypoint
warp/
  ...                     # compatibility wrappers that re-export larc
configs/
  larc_*.yaml             # training, caching, evaluation, and smoke configs
scripts/
  link_artifacts.sh       # optional helper for linking external artifacts
external/lewm_src/
  ...                     # minimal LeWorldModel compatibility source

Required Artifacts

Large data and model artifacts are not included in this repository. Prepare them yourself and place or link them to these paths:

external/lewm_object.ckpt
data/pusht_expert_train.h5
data/pusht_frame_latents.pt
checkpoints/larc_chunk_backbone_latent/shared_policy_backbone.pt        # needed for WM-FT smoke/fine-tuning
checkpoints/larc_chunk_backbone_latent/action_stats.pt

You can either copy files into the repository paths above or create symlinks from wherever your storage lives.

Generic Linux example:

mkdir -p data external checkpoints
ln -s /path/to/lewm_object.ckpt external/lewm_object.ckpt
ln -s /path/to/pusht_expert_train.h5 data/pusht_expert_train.h5
ln -s /path/to/pusht_frame_latents.pt data/pusht_frame_latents.pt

If all artifacts are in one directory, the helper script can create the standard links:

bash scripts/link_artifacts.sh /path/to/LARC_data

Installation

Use Python 3.10 or newer. Install a PyTorch build appropriate for your GPU and CUDA driver first. The experiments here were smoke-tested on AutoDL with an RTX 5090 and torch 2.8.0+cu128.

conda create -n larc python=3.10 -y
conda activate larc
pip install -r requirements-autodl.txt
pip install -e .

If you use the AutoDL image used for these experiments, torch, stable_pretraining, and stable_worldmodel may already be installed. The verified dataset package version was datasets==2.21.0.

Smoke Test

First check imports and command entrypoints:

python -c "import larc, warp; from larc.build import _patch_transformers_config; print('imports ok')"
python -m larc.train_coral --help
python -m larc.eval_coral_pusht --help
python -m warp.train_coral --help

Then verify the LeWorldModel checkpoint can be loaded:

python - <<'PY'
from larc.config import load_config
from larc.build import load_world_model

cfg = load_config("configs/larc_cache_pusht_frames.yaml")
wm = load_world_model(cfg["world_model"])
print("world model loaded:", type(wm).__name__)
PY

If data/pusht_frame_latents.pt and checkpoints/larc_chunk_backbone_latent/ are available, run a minimal WM-FT training smoke test:

python -m larc.train_coral --config configs/smoke_larc_pretrain_chunk_latent_wm_ft.yaml

A successful smoke run prints bc_loss, wm_loss, and smooth_loss, and writes outputs under checkpoints/smoke_larc_chunk_wm_ft/.

Reproducing LARC

Cache frame latents:

python -m larc.frame_cache \
  --config configs/larc_cache_pusht_frames.yaml \
  --output data/pusht_frame_latents.pt

Train latent IDM:

python -m larc.train_coral --config configs/larc_pretrain_idm_latent.yaml

Train latent action chunking:

python -m larc.train_coral --config configs/larc_pretrain_chunk_latent.yaml

Fine-tune the chunk policy with frozen world-model rollout consistency:

python -m larc.train_coral --config configs/larc_pretrain_chunk_latent_wm_ft.yaml

Evaluate the main PushT settings:

python -m larc.eval_coral_pusht --config configs/eval_larc_idm_backbone_h25_remaining_pusht.yaml
python -m larc.eval_coral_pusht --config configs/eval_larc_idm_backbone_h35_remaining_pusht.yaml
python -m larc.eval_coral_pusht --config configs/eval_larc_idm_backbone_h50_remaining_pusht.yaml

python -m larc.eval_coral_pusht --config configs/eval_larc_chunk_backbone_h25_remaining_pusht.yaml
python -m larc.eval_coral_pusht --config configs/eval_larc_chunk_backbone_h35_remaining_pusht.yaml
python -m larc.eval_coral_pusht --config configs/eval_larc_chunk_backbone_h50_remaining_pusht.yaml

python -m larc.eval_coral_pusht --config configs/eval_larc_chunk_backbone_wm_ft_h25_remaining_pusht.yaml
python -m larc.eval_coral_pusht --config configs/eval_larc_chunk_backbone_wm_ft_h35_remaining_pusht.yaml
python -m larc.eval_coral_pusht --config configs/eval_larc_chunk_backbone_wm_ft_h50_remaining_pusht.yaml

AutoDL Notes

AutoDL is not required, but the original experiments were run there. A typical AutoDL layout is:

/root/LARC                         # repository
/root/autodl-tmp/LARC_data/         # large artifacts

Link artifacts with:

cd /root/LARC
bash scripts/link_artifacts.sh /root/autodl-tmp/LARC_data

See AUTODL_STEPS.md for a more detailed AutoDL-oriented walkthrough.

Compatibility

Existing commands such as python -m warp.train_coral ... still work through thin wrappers, but new scripts and documentation should use python -m larc....

About

Latent Action Chunking with World-Model Rollout Consistency

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors