Scaling Short-Term Memory of Visuomotor Policies for Long-Horizon Tasks

Website | ReMemBench | Dataset | arXiv

Rutav Shah^1,2, Rajat Kumar Jenamani¹, Xiaohan Zhang¹, Lingfeng Sun¹, Roberto Martín-Martín², Yuke Zhu², Deva Ramanan¹, Karl Schmeckpeper¹

¹Robotics and AI Institute, ²The University of Texas at Austin

Abstract

We present PRISM, a transformer-based architecture for visuomotor policies that effectively uses short-term memory via two key components: (i) gated attention, which selectively filters retrieved information to suppress irrelevant details, and (ii) a hierarchical architecture that first compresses local information into compact tokens and then integrates them to capture temporally extended dependencies. Together, these mechanisms enable PRISM to scale short-term memory in visuomotor policies, addressing long-horizon tasks that require recalling information over extended time period.

Architecture

Installation

1. Set up environment variables (required before all subsequent steps):

export PRISM_DATAROOT=/path/to/prism/data  # Path where the data will be stored
export EXP_STORAGE_BASE_DIR=/path/to/experiments  # Path where training experiments will be saved

2. Clone required repositories:

# Clone ReMemBench repository
git clone --branch=latest git@github.com:ShahRutav/ReMemBench.git ReMemBench

# Clone RoboSuite repository
git clone --branch=abs_robot git@github.com:ShahRutav/robosuite.git robosuite

# Clone the PRISM repository
git clone --branch=cleanupv2 git@github.com:xzhang-bdai/memory-visuomotor-policies.git memory-visuomotor-policies

3. Create conda environment:

conda create -n prism python=3.10 pip -y
conda activate prism

4. Install PyTorch with CUDA support:

pip install --index-url https://download.pytorch.org/whl/cu124 \
    torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1

5. Install NVIDIA DALI:

pip install --extra-index-url https://pypi.nvidia.com nvidia-dali-cuda120

6. Install repositories in editable mode:

pip install -e ReMemBench
pip install -e robosuite
# Pin mujoco after robosuite to satisfy robocasa's strict version requirement
pip install mujoco==3.2.6
# Install main repository
pip install -e memory-visuomotor-policies

For evaluation, also follow the installation instructions from the ReMemBench repository.

7. Download kitchen assets (required for evaluation):

python ReMemBench/robocasa/scripts/download_kitchen_assets.py

8. Download the pretrained vision encoder:

# Download CrossMAE vision encoder
mkdir -p $PRISM_DATAROOT/crossmae_rtx
wget https://huggingface.co/mlfu7/ICRT/resolve/main/crossmae_rtx/cross-mae-rtx-vitb.pth -O $PRISM_DATAROOT/crossmae_rtx/cross-mae-rtx-vitb.pth

9. Download the dataset for training

huggingface-cli download Rutav/ReMemBench-Dataset \
  --repo-type dataset \
  --local-dir $PRISM_DATAROOT/memory \
  --local-dir-use-symlinks False

Training

To run PRISM experiments with the default configuration, use the following command:

python run_trainer.py -bs 1 -ng 8 -mc libero_1x.json -dc task_remembench_four.json -s 2 --gating-flag block_sigmoid_g1 -ds 8 -sl 2048 --compile-model --wandb-project-name prism_exps -br 101
python run_trainer.py -bs 1 -ng 8 -mc libero_1x.json -dc task_remembench_eight.json -s 2 --gating-flag block_sigmoid_g1 -ds 8 -sl 2048 --compile-model --wandb-project-name prism_exps -br 101

We recommend using 8xH100 GPUs and training for approximately two days to achieve solid results. You may be able to use a smaller batch size per gradient step (this was not extensively optimized) and obtain good performance with less compute.

Evaluation

CUDA_VISIBLE_DEVICES=0 MUJOCO_ELG_DEVICE_ID=0 MUJOCO_GL=egl OMP_NUM_THREADS=1 MPI_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python scripts/eval_casa.py \
    --ckpt_path /path-to-directary/checkpoint-100.pth \
    --vision_encoder_path $PRISM_DATAROOT/crossmae_rtx/cross-mae-rtx-vitb.pth \
    --task_name TASK_NAME \
    --robots PandaOmron --n_eval 25 --seed 0 

Evaluating Pretrained Checkpoints

1. Download the pre-trained model using the following command (~11 GB)

from huggingface_hub import snapshot_download
snapshot_download("Rutav/PRISM-ReMemBench-Four", repo_type="model", local_dir="PRISM-ReMemBench-Four")

hf download Rutav/PRISM-ReMemBench-Four --repo-type model --local-dir PRISM-ReMemBench-Four

2. Launch evaluations using the pretrained model

# Example for the task MemFruitInSinkRightFar
CUDA_VISIBLE_DEVICES=0 MUJOCO_ELG_DEVICE_ID=0 MUJOCO_GL=egl OMP_NUM_THREADS=1 MPI_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python scripts/eval_casa.py \
    --ckpt_path /path-to-directary/checkpoint-prism-50.pth \
    --vision_encoder_path $PRISM_DATAROOT/crossmae_rtx/cross-mae-rtx-vitb.pth \
    --task_name MemFruitInSinkRightFar \
    --robots PandaOmron --n_eval 25 --seed 0

Citation

If you find this work useful, please cite:

TODO: Add BibTex citation

Acknowledgments

We thank the authors of ICRT for their codebase, which this work builds upon.