Official PyTorch implementation of LGSR, our language-guided state representation for policy learning with world models.
If you find our codebase useful for your research, please cite our paper as:
@article{zhang2026improving,
title={Improving Policy Learning via Language-Guided State Representation in World Models},
author={Zhang, Lixuan and Kan, Meina and Shan, Shiguang and Chen, Xilin},
journal={IEEE Robotics and Automation Letters},
year={2026}
}Multi-task world models must represent many aspects of a scene in a limited latent space, and can therefore spend capacity on task-irrelevant information. LGSR uses task language as a prior to focus the learned representation on task-relevant state. In addition to reconstructing future observations, the world model predicts normalized time-to-go as an auxiliary proxy-reward objective, encouraging its latent state to capture task progress. We study the resulting representation on CALVIN language-conditioned manipulation.
Run the following steps in order from the repository root.
The released code is tested on Ubuntu 22.04.5 LTS with Python 3.10.16, PyTorch 2.7.0, and CUDA 12.6.
Create the Conda environment:
conda create -n lgsr python=3.10.16 pip -y
conda activate lgsr
python -m pip install --upgrade pipInstall CALVIN:
mkdir -p third_party
git clone --recurse-submodules https://github.com/mees/calvin.git third_party/calvin
git -C third_party/calvin checkout fa03f01f19c65920e18cf37398a9ce859274af76
git -C third_party/calvin submodule update --init --recursive
python -m pip install setuptools==57.5.0 wheel cmake==3.18.4.post1
python -m pip install --no-build-isolation pyhash==0.9.3
python -m pip install --no-build-isolation MulticoreTSNE==0.1
python -m pip install -e third_party/calvin/calvin_env/tacto
python -m pip install -e third_party/calvin/calvin_env
python -m pip install --no-deps -e third_party/calvin/calvin_modelsInstall PyTorch and the remaining dependencies:
python -m pip install torch==2.7.0 torchvision==0.22.0 \
--index-url https://download.pytorch.org/whl/cu126
python -m pip install -r requirements.txtInstall the tested OpenAI CLIP revision:
git clone https://github.com/openai/CLIP.git third_party/CLIP
git -C third_party/CLIP checkout dcba3cb2e2827b402d2701e7e1c7d9fed8a20ef1
python -m pip install --no-build-isolation third_party/CLIPVerify the installation after the final PyTorch stack is in place:
python -c "import torch, torchvision, xformers, clip, calvin_env, calvin_agent, networkx, urdfpy; print(torch.__version__, torchvision.__version__, xformers.__version__, networkx.__version__)"
python train.py --help
python eval_calvin.py --helpThe printed core versions should be torch==2.7.0, torchvision==0.22.0, and xformers==0.0.30.
Download the CALVIN ABC-D split with CALVIN's data script:
export LGSR_ROOT="$(pwd)"
export CALVIN_ROOT="$LGSR_ROOT/third_party/calvin"
cd "$CALVIN_ROOT/dataset"
bash download_data.sh ABC
cd "$LGSR_ROOT"The training and validation data will be stored under:
third_party/calvin/dataset/task_ABC_D/training/
third_party/calvin/dataset/task_ABC_D/validation/
Generate the proxy-reward annotation file used by the main LGSR experiment:
python -m scripts.reward_anno \
--annotations "$CALVIN_ROOT/dataset/task_ABC_D/training/lang_annotations/auto_lang_ann.npy" \
--output "$CALVIN_ROOT/dataset/proxy_reward/task_ABC_D/training/scaled_step_dis.npz"This command creates the proxy-reward file:
third_party/calvin/dataset/proxy_reward/task_ABC_D/training/scaled_step_dis.npz
The file stores the reward sequence for each annotated instruction.
The four partial-data subsets are provided in annotations/partial_data/. The following command reproduces one subset at a chosen output path:
python -m scripts.get_partial_data \
--annotations "$CALVIN_ROOT/dataset/task_ABC_D/training/lang_annotations/auto_lang_ann.npy" \
--ratio 0.4 --seed 42 \
--output /path/to/generated/partial_task_data_0.40.jsonChecked instruction-level word-type annotations are provided in annotations/word_type/. word_type_anno expands the checked training labels into dataset order. eval_word_type_anno converts the checked labels for the 34 CALVIN evaluation instructions into the JSON mapping used by word-mask evaluation:
mkdir -p generated
python -m scripts.word_type_anno \
--annotations "$CALVIN_ROOT/dataset/task_ABC_D/training/lang_annotations/auto_lang_ann.npy" \
--checked annotations/word_type/calvin_abc_d_training.jsonl \
--output "$CALVIN_ROOT/dataset/word_type/task_ABC_D/training/word_type_ann.jsonl"
python -m scripts.eval_word_type_anno \
--annotations-yaml "$CALVIN_ROOT/calvin_models/conf/annotations/new_playtable_validation.yaml" \
--checked annotations/word_type/calvin_abc_d_evaluation.jsonl \
--output "$LGSR_ROOT/generated/new_playtable_validation_word_types.json"These commands create the word-type files:
third_party/calvin/dataset/word_type/task_ABC_D/training/word_type_ann.jsonl
generated/new_playtable_validation_word_types.json
The training file stores one GCASLO label sequence for each language annotation. The evaluation file maps the 34 CALVIN evaluation instructions to their label sequences.
LGSR uses the official MAE ViT-Base checkpoint and OpenAI CLIP ViT-B/32 checkpoint. Download them locally:
mkdir -p checkpoints/mae checkpoints/clip
python -c "from urllib.request import urlretrieve; urlretrieve('https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth', 'checkpoints/mae/mae_pretrain_vit_base.pth')"
python -c "import clip; clip.load('ViT-B/32', device='cpu', download_root='checkpoints/clip')"Set the paths used by every launcher:
export LGSR_ROOT="$(pwd)"
export CALVIN_ROOT="$LGSR_ROOT/third_party/calvin"
export CALVIN_DATA="$CALVIN_ROOT/dataset/task_ABC_D"
export CALVIN_CONFIG="$CALVIN_ROOT/calvin_models/conf"
export MAE_CHECKPOINT="$LGSR_ROOT/checkpoints/mae/mae_pretrain_vit_base.pth"
export CLIP_CHECKPOINT="$LGSR_ROOT/checkpoints/clip/ViT-B-32.pt"
export OUTPUT_DIR="$LGSR_ROOT/outputs"
export NPROC_PER_NODE=8The released configurations comprise one proxy-supervised run, one full-data no-proxy run, four data ratios, and six instruction-component masks:
bash scripts/CALVIN_ABC_D/train_lgsr_proxy.sh
bash scripts/CALVIN_ABC_D/train_lgsr_no_proxy.sh
bash scripts/CALVIN_ABC_D/train_lgsr_partial.sh 0.1
bash scripts/CALVIN_ABC_D/train_lgsr_partial.sh 0.2
bash scripts/CALVIN_ABC_D/train_lgsr_partial.sh 0.4
bash scripts/CALVIN_ABC_D/train_lgsr_partial.sh 0.7
bash scripts/CALVIN_ABC_D/train_lgsr_word_mask.sh G
bash scripts/CALVIN_ABC_D/train_lgsr_word_mask.sh C
bash scripts/CALVIN_ABC_D/train_lgsr_word_mask.sh A
bash scripts/CALVIN_ABC_D/train_lgsr_word_mask.sh S
bash scripts/CALVIN_ABC_D/train_lgsr_word_mask.sh L
bash scripts/CALVIN_ABC_D/train_lgsr_word_mask.sh OThe four train_lgsr_partial.sh runs with ratios 0.1, 0.2, 0.4, and 0.7 are the partial-data experiments. Generate the proxy-reward annotation file before train_lgsr_proxy.sh, and generate the word-type annotation file before train_lgsr_word_mask.sh. The released setting uses 8 local GPUs. Set NPROC_PER_NODE to the number of local GPUs available on your system.
Set the model path and evaluation output directory:
export CHECKPOINT_PATH=/path/to/selected_checkpoint.pth
export EVAL_OUTPUT_ROOT="$LGSR_ROOT/eval_logs"Evaluate a proxy-supervised model with:
bash scripts/CALVIN_ABC_D/eval.sh "$CHECKPOINT_PATH" --resampler_prpEvaluate a no-proxy model with:
bash scripts/CALVIN_ABC_D/eval.sh "$CHECKPOINT_PATH"Evaluate the G word-mask model with:
bash scripts/CALVIN_ABC_D/eval.sh "$CHECKPOINT_PATH" \
--word_type_mask G \
--eval_word_type_annotations "$LGSR_ROOT/generated/new_playtable_validation_word_types.json"Use C, A, S, L, or O in place of G for the other masking runs.
This codebase is built upon Seer and flamingo-pytorch. We thank the authors for sharing their work.