Skip to content

Agentic-Intelligence-Lab/HiMem-WAM

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

5 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

HiMem-WAM:Hierarchical Memory-Gated World Action Models for Robotic Manipulation.

arXiv

Overview

HiMem-WAM is a Hierarchical Memory-Gated World Action Model for robotic manipulation. It combines motion-centric latent actions, high-level skill latents, and boundary-triggered memory updates for long-horizon manipulation. During inference, it predicts actions causally from current observations, language instructions, proprioception, and memory states, without requiring test-time future video generation or optical-flow estimation.

See the paper for the full method and benchmark results.

Highlights

  • Hierarchical latent action and skill modeling
  • Boundary-triggered memory updates
  • Causal policy execution without test-time future-video generation
  • LIBERO / LIBERO-Plus evaluation support
  • Clean checkpoint and dataset statistics based inference

Repository Structure

HiMem-WAM/
├── himem_wam/              # Core package: models, policies, memory, runtime, utilities
├── configs/                # Hydra and YAML configs for data, model, tasks, and evaluation
├── scripts/
│   ├── infer_libero.py     # LIBERO / LIBERO-Plus inference entry point
│   └── train.py            # Training and fine-tuning entry point
├── experiments/libero/     # LIBERO evaluation helpers and runner utilities
├── tests/                  # Lightweight package and interface checks
├── pyproject.toml
└── LICENSE

Installation

conda create -n himem-wam python=3.10 -y
conda activate himem-wam

python -m pip install -U pip setuptools wheel

# Install PyTorch according to your CUDA version.
python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128

python -m pip install -e .

Install MuJoCo, robosuite, LIBERO, and LIBERO-Plus according to their own project instructions. Keep simulator assets, datasets, checkpoints, and result directories outside version control.

LIBERO and LIBERO-Plus Environments

Official LIBERO and LIBERO-Plus both provide a libero Python package, so use separate environments to avoid package shadowing:

  • himem-wam-libero: install official LIBERO
  • himem-wam-libero-plus: install LIBERO-Plus

Check the active package before running evaluation:

python - <<'PY'
import libero
from libero.libero import benchmark
print("libero path:", libero.__file__)
print("benchmark keys:", sorted(benchmark.get_benchmark_dict().keys()))
PY

The official LIBERO environment should not point to a LIBERO-Plus source tree. LIBERO-Plus task counts are usually larger, so start evaluation with a specific --task.

Checkpoints

export CKPT=/path/to/himem-wam_libero.pt
export STATS=/path/to/himem-wam_libero_dataset_stats.json

Checkpoints and dataset statistics are not stored in this repository. The checkpoint and dataset statistics must match the training data, action normalization, camera layout, proprioception keys, and task distribution.

For the Wan backend, also provide the model asset roots through flags or environment variables:

export HIMEM_WAM_MODEL_BASE_PATH=/path/to/himem-wam_model_base
export HIMEM_WAM_PRETRAIN_MODEL_BASE_PATH=/path/to/himem-wam_pretrain_model_base

Do not commit checkpoints, dataset statistics, simulator assets, rollout videos, or generated outputs.

Evaluation on Official LIBERO

conda activate himem-wam-libero

export CKPT=/path/to/himem-wam_libero.pt
export STATS=/path/to/himem-wam_libero_dataset_stats.json
export CUDA_VISIBLE_DEVICES=0

python scripts/infer_libero.py \
  --backend wan \
  --config configs/libero/infer_libero.yaml \
  --checkpoint "$CKPT" \
  --dataset-stats "$STATS" \
  --suite libero_spatial \
  --task 0 \
  --num-rollouts 50 \
  --device cuda:0 \
  --save-dir outputs/libero_spatial_task0 \
  --max-steps 500

--task is optional; without it, the runner may evaluate the whole suite. --num-rollouts controls the rollout count per task. A very short horizon is only useful for checking loading and environment interaction, while reported success rates should use 500 steps or the benchmark default.

Evaluation on LIBERO-Plus

conda activate himem-wam-libero-plus

export CKPT=/path/to/himem-wam_libero.pt
export STATS=/path/to/himem-wam_libero_dataset_stats.json
export CUDA_VISIBLE_DEVICES=0

python scripts/infer_libero.py \
  --backend wan \
  --config configs/libero/infer_libero.yaml \
  --checkpoint "$CKPT" \
  --dataset-stats "$STATS" \
  --hydra-task libero_plus/2camera224 \
  --suite libero_spatial \
  --task 0 \
  --num-rollouts 50 \
  --device cuda:0 \
  --save-dir outputs/libero_plus_spatial_task0 \
  --max-steps 500

LIBERO-Plus may contain many more tasks than official LIBERO for the same suite name. Start with a single task, then expand to multi-task or full-suite evaluation. Do not use one rollout or a very short horizon as a performance conclusion.

Multi-task Evaluation

for TASK_ID in 0 1 2 10 50 100; do
  python scripts/infer_libero.py \
    --backend wan \
    --config configs/libero/infer_libero.yaml \
    --checkpoint "$CKPT" \
    --dataset-stats "$STATS" \
    --suite libero_spatial \
    --task "$TASK_ID" \
    --num-rollouts 10 \
    --device cuda:0 \
    --save-dir "outputs/libero_task_${TASK_ID}" \
    --max-steps 500
done

Use task IDs that exist in the active LIBERO or LIBERO-Plus installation.

Training

conda activate himem-wam-libero-plus

export CKPT=/path/to/himem-wam_libero.pt
export STATS=/path/to/himem-wam_libero_dataset_stats.json

python scripts/train.py \
  task=libero_plus/2camera224 \
  resume="$CKPT" \
  data.dataset_stats_path="$STATS" \
  output_dir=outputs/train_libero_plus

task=libero_plus/2camera224 is the current public config name for LIBERO-Plus training. Use task=libero/2camera224 for official LIBERO and task=agilex_piper/4camera448 for the AgileX Piper config when the matching dataset and statistics are available.

Full training recipes and released checkpoints will be updated progressively. Weight-only loading does not restore optimizer or scheduler state unless a full trainer state is provided. Do not commit checkpoints or outputs.

Development Checks

python -m compileall himem_wam scripts experiments
python -m pytest tests -q

These checks validate importability, configuration loading, checkpoint metadata, and lightweight policy interfaces. They do not run simulator benchmarks.

Citation

@misc{sun2026himemwam,
  title={HiMem-WAM: Hierarchical Memory-Gated World Action Models for Robotic Manipulation},
  author={Xiaoquan Sun and others},
  year={2026},
  eprint={2606.10363},
  archivePrefix={arXiv},
  primaryClass={cs.RO},
  url={https://arxiv.org/abs/2606.10363}
}

License

HiMem-WAM is released under the MIT License. See LICENSE for the project license.

About

Hierarchical memory-gated world-action models for long-horizon robotic manipulation.

Resources

License

Stars

6 stars

Watchers

1 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors