Skip to content

simchowitzlabpublic/much-ado-about-noising

Repository files navigation

Minimum Flow Policies Implementation

License: MIT Python 3.12+ PyTorch arXiv


Project WebsiteDocumentationDatasetCheckpoints

This repository contains the code for the paper "Much Ado About Noising: Dispelling the Myths of Generative Robotic Control". This repository is a PyTorch-based framework for behavior cloning with flow matching and related generative models, incorporating best practices from diffusion model training.

Features

  • 🧩 Clean & Modular: Composable components for losses, samplers, networks, and encoders
  • Fast: Optimized with torch.compile and CUDA graphs for maximum throughput
  • 📊 Best Practices: EMA, warmup scheduling, auto-resume, and proven training techniques
  • 🎯 Diverse Algorithms: Support for flow matching, consistency models, shortcut models, and regression
  • 🤖 Robot-Ready: Pre-configured for Robomimic, Kitchen, and PushT tasks

Documentation

Please refer to the documentation for more details.

Installation

uv sync
# install for development
uv sync --extra dev

Quick Start

Training

# on headless machine
export MUJOCO_GL=egl
# on ubuntu machine without mujoco installed
sudo apt-get install -y libglew-dev libosmesa6-dev patchelf
# Train Robomimic (state observations)
uv run examples/train_robomimic.py \
    task=lift_ph_state \
    network=chiunet \
    optimization.loss_type=flow \
    log.wandb_mode=online

# Train Robomimic (image observations)
uv run examples/train_robomimic.py \
    task=lift_ph_image \
    network=chiunet \
    optimization.batch_size=256

# Train Kitchen
uv run examples/train_kitchen.py task=kitchen_state

# Train PushT
uv run examples/train_pusht.py task=pusht_state

Evaluation

You can download checkpoints from Hugging Face.

# Evaluate trained model
uv run examples/train_robomimic.py \
    mode=eval \
    optimization.model_path="/path/to/checkpoint.pt"

Configuration

# Debug mode (quick test)
uv run examples/train_robomimic.py -cn exps/debug.yaml

# Override parameters
uv run examples/train_robomimic.py task.horizon=16 optimization.batch_size=512

# Multi-run (sweep multiple configs)
uv run examples/train_robomimic.py task=lift_ph_state,can_ph_state --multirun

See the Configuration Guide for more details.

Supported Training Objectives

This repository supports multiple training objectives:

  • Flow Matching (flow): Standard continuous normalizing flow
  • Regression (regression): Direct supervised learning baseline
  • MIP (mip): Minimum Iterative Policy with two-step sampling
  • TSD (tsd): Two-Stage Denoising
  • CTM (ctm): Consistency Trajectory Model
  • PSD (psd): Progressive Self-Distillation
  • LSD (lsd): Lagrangian Self-Distillation
  • ESD (esd): Euler Self-Distillation
  • MF (mf): Mean Flow

Porting MIP to Your Pipeline

If you are interested in porting MIP to your own pipeline, you can refer to examples/pi0_pytorch.diff. This file demonstrates how to modify the loss function and integrator of a standard flow policy to implement MIP.

⚠️ Note on Reproducibility > Since we have migrated to the latest versions of dependencies (e.g., robomimic) and performed significant code cleanup, the results obtained from this repository may not exactly replicate the specific numbers reported in the paper. However, the overall performance trends and conclusions remain the same. Original results is obtained with cleandiffuser robomimic environment and dataset.

Known Issues

  • CUDA graphs not supported for image-based tasks (requires static tensor shapes)
  • Kitchen tasks require MuJoCo 3.1.6: uv pip install "mujoco==3.1.6"

See Troubleshooting for more issues and solutions.

Citation

@article{pan2025adonoisingdispellingmyths,
      title={Much Ado About Noising: Dispelling the Myths of Generative Robotic Control},
      author={Chaoyi Pan and Giri Anantharaman and Nai-Chieh Huang and Claire Jin and Daniel Pfrommer and Chenyang Yuan and Frank Permenter and Guannan Qu and Nicholas Boffi and Guanya Shi and Max Simchowitz},
      year={2025},
      eprint={2512.01809},
      archivePrefix={arXiv},
      primaryClass={cs.RO},
      url={https://arxiv.org/abs/2512.01809},
}

About

A optimized PyTorch framework for behavior cloning with flow related generative models.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors