Gated Memory Policy

Yihuai Gao, Jinyun Liu, Shuang Li, Shuran Song

Stanford University

Project Website, ArXiv, Models, Datasets

This repository contains source code for gated memory policy training, simulation data collection and evaluation (Memimic & RoboMimic and Mikasa-Robo benchmarks), and real-world robot deployment with in-the-wild checkpoints.

We've organized our code as separate folders so you can easily take any component you need and plug it into your own system.

Repo	What it does
imitation-learning-policies/	Policy training and inference serving
real-env/	Real-world robot deployment (Natively support UR5, ARX5)
mujoco-env/	MuJoCo sim, data collection and evaluation for the Memimic (Ours) and RoboMimic benchmarks
mikasa-robo-env/	ManiSkill sim, data collection and evaluation for the Mikasa-Robo benchmark

Table of Contents

• Policy Training and Serving
  • Installation
  • Download Checkpoints and Datasets
  • Serve a Checkpoint
  • Rollout all Simulation Checkpoints
  • Train a Policy
    • Config System
    • Simulation
    • Real-world
  • Misc.
• Real-World Deployment
  • Python Environment
  • Hardware Setup
    • UR5 or UR5e
    • WSG50 with iPhUMI
    • ARX5 with iPhUMI
    • iPhone
    • Webcam / GoPro
    • SpaceMouse
  • System Overview
    • Architecture
    • Customized Packages
    • Config Aggregation
  • Run Experiments
• Memimic & RoboMimic Benchmark
  • Installation
  • Usage
    • Run heurisitic policies
    • Run spacemouse teleop
    • Collect heuristic data
    • Rollout a policy
    • Serve a remote environment
    • Serve a website for viewing rollout results
  • Adding a New Task
• Mikasa-Robo Benchmark
  • Quick Start
  • Evaluation
    • Local (single checkpoint)
    • Distributed (multi-checkpoint sweep)
  • Camera Resolution
  • Repo Layout

Code Acknowledgments

We are grateful to the following amazing open-sourced projects that made this work possible:

• iPhUMI and UMI for the portable data collection system.
• Diffusion Policy and RDT-1B for the policy model and training framework.
• RoboMimic, RoboSuite, Mikasa-Robo, and ManiSkill for the simulation benchmarks.

Citation

If you find this work useful, please cite:

@misc{gao2026gatedmemorypolicy,
  title         = {Gated Memory Policy},
  author        = {Yihuai Gao and Jinyun Liu and Shuang Li and Shuran Song},
  year          = {2026},
  eprint        = {2604.18933},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url           = {https://arxiv.org/abs/2604.18933},
}