Much Ado About Noising

https://simchowitzlabpublic.github.io/much-ado-about-noising
Admin
A PyTorch framework for generative robotic control and behavior cloning, implementing diverse...
Tokens:10,376
Snippets:93
Trust Score:8.3
Update:10 months ago
Show doc for...
Tokens:
Raw
### Install Dependencies and Train Policy (Bash)

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/index

This example demonstrates how to install the necessary dependencies, including robomimic, and then train a flow matching policy on the Robomimic Lift task using a ChiUNet network and flow loss. It utilizes `uv sync` for dependency management and `uv run` for executing the training script.

```bash
# Install dependencies
uv sync --extras robomimic

# Train a flow matching policy on Robomimic Lift task
uv run examples/train_robomimic.py \
    task=lift_ph_state \
    network=chiunet \
    optimization.loss_type=flow \
    log.wandb_mode=online
```

--------------------------------

### Install Package with Pip

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/help/faq

Instructions for installing the project using pip, an alternative to the recommended UV package manager. This command installs the package in editable mode.

```bash
pip install -e .
```

--------------------------------

### Basic Usage - Running Training Scripts

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/getting-started/configuration

These commands demonstrate basic usage of the `uv run` command to execute training scripts with different configurations. They show how to use default settings, select specific task and network configurations, override parameters, specify loss types, handle image-based tasks, enable debug mode, and perform multi-run experiments.

```bash
uv run examples/train_robomimic.py

uv run examples/train_robomimic.py task=lift_ph_state network=chiunet

uv run examples/train_robomimic.py task.horizon=16 optimization.batch_size=256

uv run examples/train_robomimic.py optimization.loss_type=flow

uv run examples/train_robomimic.py task=lift_ph_image network=chiunet

uv run examples/train_robomimic.py -cn exps/debug.yaml

uv run examples/train_robomimic.py task=lift_ph_state,can_ph_state --multirun
```

--------------------------------

### Euler Self-Distillation Training and Inference (Python)

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/getting-started/design

Applies Euler's method to self-distillation for training a flow map. The distillation loss incorporates the partial derivative of the flow map with respect to the start time and the velocity field at that time.

```python
L_SD = L_b + L_D
L_b = ‖b_t(I_t) - İ_t‖²
L_D = ‖∂_s Φ_{s,t}(I_s) + ∇ Φ_{s,t}(I_s) · b_s(I_s)‖²
```

```python
â = Φ_{0,1}(z),  where z ~ N(0, I)
```

--------------------------------

### Create Python Test Script for Task Configurations

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

Generates a Python script (`examples/test_<task_name>.py`) to verify all task configurations. It defines tasks, test steps, batch size, and functions to test individual configurations and all configurations.

```python
TASKS = {
    "task_name": ["state", "image", "keypoint"],
}

TEST_STEPS = 10  # Quick test
BATCH_SIZE = 4

def test_task(task_name, obs_type):
    """Test a single configuration."""
    cmd = [
        "uv", "run", "python",
        "examples/train_task.py",
        "-cn", "exps/debug.yaml",
        f"task={task_name}_{obs_type}",
        f"optimization.gradient_steps={TEST_STEPS}",
        # ... other overrides
    ]
    # Run and verify

def main():
    """Test all configurations."""
    # Run tests and report results

```

--------------------------------

### Run Training Tests in Bash

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

Bash commands to test the training process with different configurations. It includes options for using a local dataset with a debug config or a HuggingFace dataset, and running the full test suite.

```bash
# Test with debug config using local dataset
uv run python examples/train_<task_name>.py -cn exps/debug.yaml task=<task_name>_state +task.dataset_path=data/<task_name>/<dataset_dir>

# Test with HuggingFace dataset (if uploaded)
uv run python examples/train_<task_name>.py -cn exps/debug.yaml task=<task_name>_state

# Run full test suite
uv run python examples/test_<task_name>.py
```

--------------------------------

### Quick Dataset Test in Bash

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

A bash command to quickly test the dataset loader for a specific task. It instantiates the dataset, prints the number of samples, and checks the shapes of the state and action tensors.

```bash
# Quick test to verify dataset loading
uv run python -c "from mip.datasets.<task_name>_dataset import <Task>StateDataset; \ 
ds = <Task>StateDataset('data/<task_name>/<dataset_dir>', horizon=16, pad_before=1, pad_after=7); \ 
print(f'✅ Dataset loaded: {len(ds)} samples'); \ 
item = ds[0]; \ 
print(f'✅ State shape: {item[\"obs\"][\"state\"]}.shape'); \ 
print(f'✅ Action shape: {item[\"action\"]}.shape')"
```

--------------------------------

### Download and Process Dataset using Bash

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

Provides bash commands to download and process datasets for a given task. It includes an option to skip uploading to HuggingFace or to upload after processing.

```bash
# Download dataset (will download to data/<task_name>/)
uv run python examples/process_<task_name>_dataset.py --skip_upload

# Or upload to HuggingFace (requires authentication)
uv run python examples/process_<task_name>_dataset.py
```

--------------------------------

### Speed Up Training with torch.compile

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/help/faq

Enables torch.compile to speed up training. This optimization can significantly reduce training times by compiling PyTorch code.

```bash
uv run examples/train_robomimic.py optimization.use_compile=true
```

--------------------------------

### State Observation Processing in Python

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

Demonstrates how to process state observations in Python, both within the dataset creation and during evaluation. It shows how to structure the observation dictionary and normalize state tensors.

```python
# In dataset
obs_dict = {"state": state_tensor}

# In eval
obs = dataset.normalizer["obs"]["state"].normalize(obs)
obs = torch.tensor(obs, device=device)

```

--------------------------------

### Task Configuration for MIP Environments in YAML

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

This YAML configuration defines parameters for a task environment within the MIP project. It includes settings for environment identification, observation type, action space, dataset paths, environment parameters like `max_episode_steps` and `num_envs`, and dimensionalities for observations and actions. It supports different observation types and dataset sources.

```yaml
# @package task
defaults:
  - _self_

# Task identification
env_name: "task_name"
obs_type: "state"  # or "image", "keypoint", etc.
abs_action: false  # true for absolute actions, false for delta actions

# Dataset configuration
dataset_path: "path/to/dataset"
# OR for HuggingFace:
# dataset_repo: "ChaoyiPan/mip-dataset"
# dataset_filename: "task/variant/dataset.hdf5"

# Environment parameters
max_episode_steps: 400

# Task parameters
obs_steps: 2        # Number of observation steps to stack
act_steps: 8        # Number of action steps to execute
horizon: 16         # Prediction horizon
num_envs: 1         # Number of parallel environments
save_video: false   # Whether to save videos
val_dataset_percentage: 0.0  # Validation split percentage

# Dimensions
act_dim: 10         # Action dimension
obs_dim: 20         # Observation dimension
```

--------------------------------

### Resume MIP Training with Auto-Resume

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/help/faq

Enables auto-resume functionality for MIP training, allowing training to automatically resume from the last checkpoint if it exists. This is useful for long training sessions.

```bash
# Training will automatically resume from last checkpoint if it exists
uv run examples/train_robomimic.py task=lift_ph_state optimization.auto_resume=true
```

--------------------------------

### Helper Function to Create Single Environment Instance in Python

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

This Python function defines a 'thunk' that creates a single environment instance. It includes steps for creating the base environment, applying observation wrappers (like `ImageWrapper`), adding video recording (`VideoRecordingWrapper`), and incorporating a `MultiStepWrapper` for observation stacking and action repetition. It also handles setting the random seed.

```python
def make_task_env(task_config: TaskConfig, idx, render=False, seed=None):
    """Create a single environment instance."""
    def thunk():
        # 1. Create base environment (import locally to avoid dependency issues)
        env = create_base_env(task_config)

        # 2. Add observation wrappers if needed
        if task_config.obs_type == "image":
            env = ImageWrapper(env)

        # 3. Add video recording wrapper
        video_recoder = VideoRecorder.create_h264(...) # Assuming VideoRecorder is imported
        env = VideoRecordingWrapper(env, video_recoder, ...) # Assuming VideoRecordingWrapper is imported

        # 4. Add multi-step wrapper (important!)
        env = MultiStepWrapper(
            env,
            n_obs_steps=task_config.obs_steps,
            n_action_steps=task_config.act_steps,
            max_episode_steps=task_config.max_episode_steps,
        ) # Assuming MultiStepWrapper is imported

        # 5. Set seed
        if seed is not None:
            env.seed(seed + idx)

        return env
    return thunk
```

--------------------------------

### Define Dataset Sample Structure in Python

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

Specifies the expected output structure for the `__getitem__` method in dataset classes. The returned data should be a dictionary containing observations ('obs') and actions ('action'). Observations can be a dictionary (for image data) or a tensor (for state data), while actions are always a tensor.

```python
{
    "obs": {...},  # Observations (dict for image, tensor for state)
    "action": tensor  # Actions
}
```

--------------------------------

### Create Task Dataset Loader in Python

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

Defines the structure and components for creating a new dataset loader within the MIP framework. It requires a `make_dataset` factory function and dataset classes that inherit from `BaseDataset`, implementing methods for initialization, length, item retrieval, and normalizer access. Supports various observation types and utilizes `ReplayBuffer` and `SequenceSampler`.

```python
def make_dataset(task_config, mode="train"):
    """Create dataset based on configuration."""
    if task_config.obs_type == "state":
        return StateDataset(...)
    elif task_config.obs_type == "image":
        return ImageDataset(...)
    else:
        raise ValueError(f"Invalid observation type: {task_config.obs_type}")
```

```python
class TaskStateDataset(BaseDataset):
    def __init__(self, dataset_path, horizon=1, pad_before=0, pad_after=0):
        super().__init__()
        self.replay_buffer = ReplayBuffer.copy_from_path(dataset_path, keys=...)
        self.sampler = SequenceSampler(...)
        self.normalizer = self.get_normalizer()

    def get_normalizer(self):
        return {
            "obs": {"state": MinMaxNormalizer(...)},
            "action": MinMaxNormalizer(...)
        }

    def __len__(self):
        return len(self.sampler)

    def __getitem__(self, idx):
        sample = self.sampler.sample_sequence(idx)
        # Process and normalize sample
        return torch_data
```

--------------------------------

### Image Observation Processing in Python

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

Illustrates the processing of image observations along with agent position in Python. It covers structuring the observation dictionary and normalizing various observation types during evaluation.

```python
# In dataset
obs_dict = {
    "image": image_tensor,  # (T, C, H, W)
    "agent_pos": pos_tensor
}

# In eval
for k in obs_raw:
    obs[k] = dataset.normalizer["obs"][k].normalize(obs[k])
    obs[k] = torch.tensor(obs[k], device=device)

```

--------------------------------

### Action Processing in Python

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

Shows how to process actions in Python, including selecting action steps and undoing transformations for absolute actions with rotation. This is typically done within the evaluation loop.

```python
# In eval loop
start = config.task.obs_steps - 1
end = start + config.task.act_steps
action = action_pred[:, start:end, :]

# For absolute actions with rotation
if config.task.abs_action:
    action = dataset.undo_transform_action(action)

```

--------------------------------

### Speed Up Training with CUDA Graphs

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/help/faq

Enables CUDA graphs to reduce CPU overhead during training, specifically for state observations. This can lead to performance improvements on compatible hardware.

```bash
uv run examples/train_robomimic.py optimization.use_cudagraphs=true
```

--------------------------------

### Run Dataset Tests using Pytest

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

Bash commands to run tests for the dataset loader using `pytest`. It allows testing all dataset tests or specific dataset classes like `Test<Task>StateDataset`.

```bash
# Test dataset loader (may skip if HuggingFace download not set up)
uv run pytest mip/datasets/<task_name>_dataset_test.py -v

# Test specific observation type
uv run pytest mip/datasets/<task_name>_dataset_test.py::Test<Task>StateDataset -v
```

--------------------------------

### Create Vectorized Environments in Python

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

This Python function creates vectorized environments for a given task configuration. It chooses between synchronous and asynchronous vectorization based on the number of environments and whether video saving is enabled. It requires the `gym` library and a `TaskConfig` object.

```python
import gym
def make_vec_env(task_config: TaskConfig, seed=None):
    """Create vectorized environments.

    Args:
        task_config: Task configuration
        seed: Random seed for environments

    Returns:
        Vectorized gym environment
    """
    # Choose vectorization strategy
    if task_config.num_envs == 1 or task_config.save_video:
        vnc_env_class = gym.vector.SyncVectorEnv
    else:
        vnc_env_class = gym.vector.AsyncVectorEnv

    # Create vectorized environments
    envs = vnc_env_class([
        make_task_env(task_config, idx, False, seed=seed)
        for idx in range(task_config.num_envs)
    ])
    return envs
```

--------------------------------

### Image Observation Configuration in YAML

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

This YAML snippet extends the base task configuration to include specific settings for image-based observations. It defines the structure of the `shape_meta` for actions and observations, including image keys and low-dimensional keys, and specifies parameters for the image encoder, such as the model type (`resnet18`), resizing, cropping, and normalization techniques.

```yaml
shape_meta:
  action:
    shape: [10]
  obs:
    image_key:
      shape: [3, 96, 96]
      type: rgb
    lowdim_key:
      shape: [10]
      type: low_dim

# Image encoder settings
rgb_model: resnet18
resize_shape: null
crop_shape: [84, 84]
random_crop: true
use_group_norm: true
use_seq: true
```

--------------------------------

### Two Step Denoising Training and Inference (Python)

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/getting-started/design

A simplified version of flow matching that involves two denoising steps. Training minimizes errors across both steps, and inference uses the learned policies sequentially to generate the final output.

```python
π_θ ≈ argmin_θ 𝔼[‖π_θ(o, I_0, t=0) - I_fix‖² + ‖π_θ(o, I_fix, t_fix) - a‖²]
```

```python
â_0 ← π_θ(o, z, 0)
â ← π_θ(o, t_fix * â_0 + (1 - t_fix) * z, t_fix)
```

--------------------------------

### Minimum Iterative Policy Training and Inference (Python)

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/getting-started/design

Optimized for two-step sampling by reducing stochasticity. The training objective includes minimizing prediction errors at different stages, and inference uses a simplified, iterative approach.

```python
π_mip^θ ≈ argmin_θ 𝔼[‖π_θ(o, I_0 = 0, t=0) - a‖² + ‖π_θ(o, I_t_fix, t_fix) - a‖²]
```

```python
â_mip^0 ← π_mip^θ(o, 0, t=0)
â_mip ← π_mip^θ(o, t_fix * â_mip^0, t_fix)
```

--------------------------------

### Regression Training and Inference (Python)

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/getting-started/design

Utilizes a direct regression loss to learn the conditional action mean. The training minimizes the difference between the predicted action and the actual action, while inference directly uses the learned policy to generate actions.

```python
π_θ ≈ argmin_θ 𝔼[‖π_θ(o, I_0, t=0) - a‖²]
```

```python
â ← π_θ(o, z, t=0)
```

--------------------------------

### Update __init__.py to Export make_vec_env in Python

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_task

This Python snippet shows how to update the `__init__.py` file for an environment package. It imports the `make_vec_env` function from the environment wrapper module and adds it to the `__all__` list, making it accessible when the package is imported.

```python
from mip.envs.<task_name>.<task_name>_env_wrapper import make_vec_env

__all__ = ["make_vec_env"]
```

--------------------------------

### Progressive Self-Distillation Training and Inference (Python)

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/getting-started/design

A self-distillation framework for training a flow map. It combines a standard flow matching loss with a distillation loss that encourages consistency between different time-step transformations.

```python
L_SD = L_b + L_D
L_b = ‖b_t(I_t) - İ_t‖²
L_D = ‖Φ_{s,t}(I_s) - Φ_{u,t}(Φ_{s,u}(I_s))‖²
```

```python
â = Φ_{0,1}(z),  where z ~ N(0, I)
```

--------------------------------

### Lagrangian Self-Distillation Training and Inference (Python)

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/getting-started/design

A self-distillation framework focusing on Lagrangian dynamics for training a flow map. The distillation loss is formulated based on the time derivative of the flow map and the velocity field.

```python
L_SD = L_b + L_D
L_b = ‖b_t(I_t) - İ_t‖²
L_D = ‖∂_t Φ_{s,t}(I_s) - b_t(Φ_{s,t}(I_s))‖²
```

```python
â = Φ_{0,1}(z),  where z ~ N(0, I)
```

--------------------------------

### Consistency Trajectory Model Training and Inference (Python)

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/getting-started/design

Progressively distills multi-step flow into flow map/shortcut models. The training involves ensuring consistency across different time steps, and inference directly uses the learned flow map to transform noise into data.

```python
Φ_{s,t}(I_s) ≈ Φ_{s+dt, t}(stopgrad(Φ_{s, s+dt}(I_s)))
```

```python
a = Φ_{0,1}(z),  where z ~ N(0, I)
```

--------------------------------

### Flow Matching Training and Inference (Python)

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/getting-started/design

Implements the standard flow matching loss for generative modeling. It trains a model to predict the velocity field that transforms data from one distribution to another. The inference step uses this learned velocity field to generate new data points.

```python
b_θ ≈ argmin_θ 𝔼[‖b_t(I_t | o) - İ_t‖²],  where t ~ Unif([0,1]), z ~ N(0, I)
```

```python
d/dt a_t = b_t(a_t | o)  with initial condition  a_0 = z
```

--------------------------------

### Mean Flow Definition, Training, and Inference (Python)

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/getting-started/design

Introduces Mean Flow, a self-distillation framework that defines the flow map using a mean velocity. The training involves both a flow matching loss and a distillation loss based on the divergence theorem.

```python
Φ_{s,t}(I_s) = I_s + (t - s) * v̄_{s,t}(I_s)
```

```python
L_SD = L_b + L_D
L_b = ‖b_t(I_t) - İ_t‖²
L_D = ‖∂_s Φ_{s,t}(I_s) + stopgrad(∇ Φ_{s,t}(I_s) · İ_s)‖²
```

```python
â = Φ_{0,1}(z),  where z ~ N(0, I)
```

--------------------------------

### Define Sampler Function in Python

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_method

Defines the signature for a sampler function in `mip/samplers.py`. Samplers implement inference-time sampling strategies corresponding to loss types. They take configuration, flow map, encoder, initial noise, and observations as input, returning a predicted action trajectory.

```python
def sampler(
    config: OptimizationConfig,
    flow_map: FlowMap,
    encoder: BaseEncoder,
    act_0: torch.Tensor,  # Initial noise
    obs: torch.Tensor,    # Observations
) -> torch.Tensor:
    """
    Returns:
        action: Predicted action trajectory
    """

```

--------------------------------

### Define Loss Function in Python

Source: https://simchowitzlabpublic.github.io/much-ado-about-noising/development/add_new_method

Defines the signature for a loss function in `mip/losses.py`. Loss functions adhere to a consistent signature, taking configuration, flow map, encoder, interpolant, action, observations, and delta_t as input. They return a scalar loss tensor and a dictionary of additional metrics for logging.

```python
def loss_fn(
    config: OptimizationConfig,
    flow_map: FlowMap,
    encoder: BaseEncoder,
    interp: Interpolant,
    act: torch.Tensor,
    obs: torch.Tensor,
    delta_t: torch.Tensor,
) -> tuple[torch.Tensor, dict]:
    """
    Returns:
        loss: Scalar loss tensor
        info: Dict with additional metrics for logging
    """

```
Tokens:
Raw