What’s new here

IMPRINT attaches a trained memory encoder to vision-language-action robot policies. The encoder is policy-agnostic: the same two-stage training recipe works with Pi0, Pi0.5, and MolmoAct2, each on a separate branch with its own dependency stack. Two encoder variants are interchangeable: a full-attention encoder that lets memory queries attend over the retained history, and a recurrent encoder that maintains a fixed-size gated delta state for lower update cost on long histories.

What it does

IMPRINT gives a VLA policy access to episode history through one of two interchangeable encoders. The full-attention encoder lets memory queries attend over the retained history directly. The recurrent encoder maintains a fixed-size gated delta state, which costs less to update as histories grow long.

Training splits into two stages. Stage 1 bootstraps the memory encoder with reconstruction objectives. Stage 2 then adapts it together with the VLA policy. Both stages use paper configurations stored under configs/paper/ and write reproducibility metadata beside the checkpoints.

The repo ships evaluation harnesses for two benchmarks: RMBench (nine manipulation tasks via RoboTwin-Mem inside EventVLA) and a patched RoboCasa 1.0.1 with added memory tasks. Policy and simulator run as separate processes with separate Python environments to avoid dependency conflicts. Published checkpoints and training datasets are hosted on Hugging Face under Amazon-FAR/imprint_models and Amazon-FAR/imprint_datasets.

Who it’s for

Robotics researchers who want to add structured episode memory to an existing VLA policy and evaluate it on standardized benchmarks. You need Linux x86_64, Python 3.11, CUDA 12.8, and four 80 GB GPUs for the paper training configurations. Evaluation alone needs two GPUs (one for the policy server, one for the simulator). The real-world training pipeline and datasets are intentionally excluded from this release.

Try it

Requirements: Linux x86_64, Python 3.11, uv, Hugging Face CLI (hf), NVIDIA driver compatible with CUDA 12.8.

Download a checkpoint only (no training data):

bash scripts/reproduce/download_release.sh \
  rmbench full_attention --model-only

Install the policy environment:

git clone https://github.com/amazon-far/imprint.git
cd imprint
export UV_CACHE_DIR=$PWD/.uv-cache
export OPENPI_DATA_HOME=$PWD/.openpi-cache
UV_PROJECT_ENVIRONMENT=$PWD/.venv uv sync --python 3.11 --frozen
source .venv/bin/activate
python -c "import jax; print(jax.devices())"

Stage the DINOv3 model (required for evaluation):

hf download facebook/dinov3-vitb16-pretrain-lvd1689m \
  --revision 5931719e67bbdb9737e363e781fb0c67687896bc \
  --local-dir artifacts/models/dinov3-vitb16
export DINOV3_MODEL_PATH=$PWD/artifacts/models/dinov3-vitb16

Run a one-trial smoke test on RMBench after setting up the simulator environment per the README:

python benchmarks/rmbench/eval/run_local.py \
  --rmbench-root "$RMBENCH_ROOT" \
  --checkpoint \
    artifacts/checkpoints/rmbench-full-attention/stage2/step-100000 \
  --dinov3-model "$DINOV3_MODEL_PATH" \
  --architecture full_attention \
  --task cover_blocks \
  --trials 1 \
  --start-seed 100000 \
  --policy-python "$PWD/.venv/bin/python" \
  --sim-python "$RMBENCH_PYTHON" \
  --policy-gpu 0 \
  --sim-gpu 1 \
  --output outputs/rmbench/cover-smoke

How mature is it

Created 2026-10-10, last pushed 2026-10-14. 8 stars, 0 forks, 1 contributor, 1 commit in the last 90 days, 4 open pull requests, 0 releases. Checkpoints and datasets are published on Hugging Face rather than as GitHub releases. Apache-2.0 license, with Gemma components under a separate LICENSE_GEMMA.txt.