### Create and Populate Zarr Replay Buffer

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Complete example demonstrating creation of a Zarr-based replay buffer with datasets for actions, images, and state observations. Registers jpeg2000 codec for compression, creates chunked storage structures, and adds episode data. Shows practical setup for managing demonstration datasets efficiently.

```Python
from diffusion_policy.common.replay_buffer import ReplayBuffer
from diffusion_policy.codecs.imagecodecs_numcodecs import register_codecs
import numpy as np
import zarr

# Register jpeg2000 codec for compression
register_codecs()

# Create new replay buffer
store = zarr.DirectoryStore('data/my_dataset.zarr')
root = zarr.group(store=store, overwrite=True)

# Define storage with compression
root.create_dataset(
    'data/action',
    shape=(0, 2),
    chunks=(100, 2),
    dtype=np.float32
)
root.create_dataset(
    'data/image',
    shape=(0, 96, 96, 3),
    chunks=(1, 96, 96, 3),
    dtype=np.uint8,
    compressor=zarr.Blosc(cname='zstd', clevel=3)
)
root.create_dataset(
    'data/state',
    shape=(0, 5),
    chunks=(100, 5),
    dtype=np.float32
)
root.create_dataset(
    'meta/episode_ends',
    shape=(0,),
    chunks=(100,),
    dtype=np.int64
)

# Add episode data
replay_buffer = ReplayBuffer.create_from_path('data/my_dataset.zarr', mode='a')

episode_data = {
    'action': np.random.randn(125, 2).astype(np.float32),
    'image': (np.random.rand(125, 96, 96, 3) * 255).astype(np.uint8),
    'state': np.random.randn(125, 5).astype(np.float32)
}

replay_buffer.add_episode(episode_data)
```

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

### DataLoader and Training Loop Example (Python)

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Shows how to create a PyTorch DataLoader from the PushTDataset and iterate through it for a training loop. It demonstrates accessing batched observations and actions. Dependencies include torch.utils.data.

```python
# Create dataset and dataloader
dataset = PushTDataset(
    zarr_path='data/pusht/pusht_cchi_v7_replay.zarr',
    horizon=16,
    n_obs_steps=2,
    n_action_steps=8
)

dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=64,
    num_workers=4,
    shuffle=True,
    pin_memory=True
)

# Training loop
for batch in dataloader:
    obs_dict = batch['obs']    # {'image': (B,To,H,W,3), 'agent_pos': (B,To,Ds)}
    actions = batch['action']  # (B, T, Da)
    # Train model...
```

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

### PushTDataset Example using SequenceSampler (Python)

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

An example of creating a PushTDataset that utilizes SequenceSampler to load and preprocess data for diffusion policy training. It handles replay buffer loading, calculates padding, splits data into training/validation sets, and creates a sampler. Dependencies include ReplayBuffer, SequenceSampler, LinearNormalizer, and torch.utils.data.

```python
# Example: Create dataset with proper padding
from diffusion_policy.common.replay_buffer import ReplayBuffer
from diffusion_policy.common.sampler import SequenceSampler
from diffusion_policy.model.common.normalizer import LinearNormalizer
import torch.utils.data as data

class PushTDataset(data.Dataset):
    def __init__(self,
        zarr_path: str,
        horizon: int = 16,
        n_obs_steps: int = 2,
        n_action_steps: int = 8,
        seed: int = 42,
        val_ratio: float = 0.02
    ):
        # Load replay buffer
        self.replay_buffer = ReplayBuffer.create_from_path(zarr_path, mode='r')

        # Calculate padding
        # pad_before allows sampling from episode start
        # pad_after allows sampling until episode end
        pad_before = n_obs_steps - 1
        pad_after = n_action_steps - 1

        # Split train/val episodes
        n_episodes = self.replay_buffer.n_episodes
        episode_indices = np.arange(n_episodes)
        np.random.seed(seed)
        np.random.shuffle(episode_indices)
        n_val = int(n_episodes * val_ratio)
        val_mask = np.zeros(n_episodes, dtype=bool)
        val_mask[episode_indices[:n_val]] = True
        train_mask = ~val_mask

        # Create sampler (train episodes only)
        self.sampler = SequenceSampler(
            replay_buffer=self.replay_buffer,
            sequence_length=horizon,
            pad_before=pad_before,
            pad_after=pad_after,
            episode_mask=train_mask
        )

        self.horizon = horizon
        self.n_obs_steps = n_obs_steps
        self.n_action_steps = n_action_steps

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

    def __getitem__(self, idx):
        # Get sequence with automatic padding
        data = self.sampler.sample_sequence(idx)

        # Extract observations and actions
        obs = {
            'image': data['image'][:self.n_obs_steps],      # (To, H, W, 3)
            'agent_pos': data['state'][:self.n_obs_steps]   # (To, Ds)
        }
        action = data['action']  # (T, Da)

        return {'obs': obs, 'action': action}
```

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

### LinearNormalizer Usage Example (Python)

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

A snippet demonstrating the complete workflow of using LinearNormalizer: collecting dataset statistics and fitting the normalizer. This is a precursor to using the normalizer in policy training. Dependencies include LinearNormalizer, PushTDataset, torch, and numpy.

```python
# Complete example: fit normalizer and use in training
from diffusion_policy.model.common.normalizer import LinearNormalizer
import torch
import numpy as np

# Collect dataset statistics
dataset = PushTDataset('data/pusht/pusht_cchi_v7_replay.zarr')


```

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

### Diffusion UNet Image Policy Training Configuration

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Complete OmegaConf configuration example for training DiffusionUnetImagePolicy with ResNet18 vision encoder, DDPM noise scheduler, and training hyperparameters. Includes policy architecture details, noise scheduler settings, observation encoder configuration with group normalization, and training parameters with EMA decay.

```Python
from diffusion_policy.workspace.train_diffusion_unet_image_workspace import TrainDiffusionUnetImageWorkspace
from omegaconf import OmegaConf

# Load configuration
cfg = OmegaConf.load('image_pusht_diffusion_policy_cnn.yaml')

# Configuration structure
config = OmegaConf.create({
    '_target_': 'diffusion_policy.workspace.train_diffusion_unet_image_workspace.TrainDiffusionUnetImageWorkspace',
    'name': 'train_diffusion_unet_image',
    'horizon': 16,
    'n_obs_steps': 2,
    'n_action_steps': 8,
    'obs_as_global_cond': True,
    'policy': {
        '_target_': 'diffusion_policy.policy.diffusion_unet_image_policy.DiffusionUnetImagePolicy',
        'noise_scheduler': {
            '_target_': 'diffusers.schedulers.scheduling_ddpm.DDPMScheduler',
            'num_train_timesteps': 100,
            'beta_schedule': 'squaredcos_cap_v2',
            'prediction_type': 'epsilon',
            'clip_sample': True
        },
        'obs_encoder': {
            '_target_': 'diffusion_policy.model.vision.multi_image_obs_encoder.MultiImageObsEncoder',
            'rgb_model': {
                '_target_': 'diffusion_policy.model.vision.model_getter.get_resnet',
                'name': 'resnet18'
            },
            'crop_shape': [76, 76],
            'random_crop': True,
            'use_group_norm': True
        },
        'diffusion_step_embed_dim': 128,
        'down_dims': [512, 1024, 2048],
        'kernel_size': 5,
        'n_groups': 8
    },
    'training': {
        'device': 'cuda:0',
        'num_epochs': 8000,
        'batch_size': 64,
        'lr': 1.0e-4,
        'use_ema': True,
        'ema_decay': 0.995,
        'rollout_every': 50,
        'checkpoint_every': 50,
        'val_every': 1
    }
})

# Initialize and train
workspace = TrainDiffusionUnetImageWorkspace(config)
workspace.run()
```

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

### BaseImagePolicy Interface and Inference

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Defines the standard interface for all vision-based policies accepting image observations and returning action sequences. Includes predict_action method with tensor shape specifications and example usage loading a trained DiffusionUnetImagePolicy checkpoint for inference with action chunking execution.

```python
import torch
from typing import Dict

class BaseImagePolicy:
    def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Predict actions from observations.

        Args:
            obs_dict: Dictionary containing:
                'image': (B, To, H, W, 3) in [0,1] float32 range
                'agent_pos': (B, To, D) low-dimensional state (optional)

        Returns:
            action_dict: Dictionary containing:
                'action': (B, Ta, Da) predicted actions

        Where:
            B: batch size
            To: observation horizon (n_obs_steps)
            Ta: action horizon (n_action_steps)
            H, W: image height/width
            D: low-dim observation dimension
            Da: action dimension
        """
        pass

# Example usage with trained policy
from diffusion_policy.policy.diffusion_unet_image_policy import DiffusionUnetImagePolicy
import torch

# Load checkpoint
checkpoint = torch.load('checkpoint.ckpt', map_location='cuda:0')
cfg = checkpoint['cfg']
policy = DiffusionUnetImagePolicy(cfg.policy)
policy.load_state_dict(checkpoint['state_dicts']['model'])
policy.eval()

# Prepare observations (2 timesteps of 96x96 RGB images)
obs_dict = {
    'image': torch.randn(1, 2, 96, 96, 3).cuda(),  # B=1, To=2
    'agent_pos': torch.randn(1, 2, 2).cuda()        # 2D position
}

# Get action predictions (8 timesteps)
with torch.no_grad():
    result = policy.predict_action(obs_dict)
    actions = result['action']  # (1, 8, 2) - 8 future actions

# Execute first n_action_steps with action chunking
for i in range(actions.shape[1]):
    env.step(actions[0, i].cpu().numpy())
```

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

### Real Robot Evaluation Script with Diffusion Policy (Python)

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

This script demonstrates how to load a trained diffusion policy and execute it on a real robot environment. It initializes the policy, shared memory, and the RealEnv, then enters a loop to collect observations, predict actions, and execute them. It requires a 'checkpoint.ckpt' file and specific camera serial numbers.

```python
from diffusion_policy.real_world.real_env import RealEnv
from diffusion_policy.policy.diffusion_unet_image_policy import DiffusionUnetImagePolicy
from diffusion_policy.shared_memory.shared_memory_manager import SharedMemoryManager
import torch
import time

# Helper function for precise sleeping (implementation not provided in snippet)
# def precise_sleep(target_time): pass 

# Load trained policy
checkpoint = torch.load('checkpoint.ckpt')
policy = DiffusionUnetImagePolicy(checkpoint['cfg'].policy)
policy.load_state_dict(checkpoint['state_dicts']['model'])
policy.set_normalizer(checkpoint['state_dicts']['normalizer'])
policy.eval().cuda()

# Create shared memory manager
shm_manager = SharedMemoryManager()
shm_manager.start()

# Create environment
env = RealEnv(
    output_dir='data/eval_real',
    robot_ip='192.168.0.204',
    frequency=10,                        # 10 Hz control
    n_obs_steps=2,
    camera_serial_numbers=['128422270679', '127122270146'],  # Two RealSense cameras
    max_pos_speed=0.25,
    max_rot_speed=0.6,
    shm_manager=shm_manager
)

# Evaluation loop
n_episodes = 10
dt = 1.0 / 10  # Control period

for episode_idx in range(n_episodes):
    # Reset robot to initial pose (assuming env.reset() is implemented)
    # env.reset()

    # Start recording
    start_time = time.time()
    env.start_episode(start_time)

    step = 0
    while step < 300:  # Max 30 seconds
        # Get observations
        obs_dict_np = env.get_obs()

        # Convert to torch and normalize images to [0, 1]
        obs_dict = {
            'image_0': torch.from_numpy(obs_dict_np['camera_0'] / 255.0).cuda(),
            'image_1': torch.from_numpy(obs_dict_np['camera_1'] / 255.0).cuda(),
            'robot_eef_pose': torch.from_numpy(obs_dict_np['robot_eef_pose']).cuda()
        }

        # Add batch dimension
        for key in obs_dict:
            obs_dict[key] = obs_dict[key].unsqueeze(0)

        # Predict actions
        with torch.no_grad():
            action_dict = policy.predict_action(obs_dict)
            actions = action_dict['action'][0].cpu().numpy()  # (n_action_steps, 6)

        # Calculate target timestamps for action chunk
        current_time = time.time()
        n_action_steps = actions.shape[0]
        timestamps = current_time + np.arange(n_action_steps) * dt

        # Execute action chunk asynchronously
        env.exec_actions(actions, timestamps)

        # Wait for next control cycle
        target_time = current_time + dt
        # precise_sleep(target_time) # Call to helper function

        step += 1

    # Stop recording and save episode
    env.end_episode()

# Cleanup
shm_manager.shutdown()

# Episodes saved to: data/eval_real/replay_buffer.zarr

```

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

### ReplayBuffer Class for Dataset Management

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Zarr-based replay buffer class for efficient dataset storage and retrieval. Provides methods to create buffers from disk paths and add episodes with heterogeneous data types (actions, images, states). Supports compression and enables entire datasets to fit in RAM for efficient training.

```Python
import numpy as np
import zarr

class ReplayBuffer:
    @classmethod
    def create_from_path(cls, zarr_path: str, mode='r'):
        """Load replay buffer from disk."""
        return cls(root=zarr.open(zarr_path, mode))

    def add_episode(self, data_dict: Dict[str, np.ndarray]):
        """
        Add episode to buffer.

        Args:
            data_dict: {'action': (T, Da), 'image': (T, H, W, 3), 'state': (T, Ds)}
        """
        pass

    @property
    def episode_ends(self) -> np.ndarray:
        """Get array of episode end indices."""
        return self.root['meta']['episode_ends'][:]
```

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

### Access Replay Buffer Data

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Demonstrates how to access actions, images, and episode end information from a replay buffer. The shapes of the extracted data are commented for clarity.

```python
actions = replay_buffer['action']  # (125, 2)
images = replay_buffer['image']    # (125, 96, 96, 3)
episode_ends = replay_buffer.episode_ends  # [125]
```

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

### Train Diffusion Policy with Hydra Configuration

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

End-to-end training script that uses Hydra configuration management to instantiate workspace, dataset, policy, and training loop. Includes command-line configuration overrides for training parameters, device specification, and output directory management using Hydra's date/time formatting.

```python
import hydra
from omegaconf import OmegaConf

@hydra.main(
    config_path='diffusion_policy/config',
    version_base=None
)
def main(cfg: OmegaConf):
    # Instantiate workspace from config
    cls = hydra.utils.get_class(cfg._target_)
    workspace = cls(cfg)
    workspace.run()

if __name__ == "__main__":
    main()
```

```bash
# Download and prepare data
mkdir data && cd data
wget https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip
unzip pusht.zip && rm pusht.zip && cd ..

# Train with custom configuration
python train.py \
  --config-name=train_diffusion_unet_image_workspace \
  task=pusht_image \
  training.seed=42 \
  training.device=cuda:0 \
  training.num_epochs=8000 \
  training.batch_size=64 \
  training.lr=1.0e-4 \
  policy.num_inference_steps=100 \
  hydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}'
```

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

### Task Configuration Schema (YAML)

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Defines the schema for configuring new tasks within the Diffusion Policy framework, adhering to the O(N+M) architecture. It specifies observation and action spaces, details the environment runner for evaluation, and configures the dataset loader for training, including parameters like replay buffer path, horizon, and data splitting.

```yaml
# diffusion_policy/config/task/pusht_image.yaml
name: pusht_image

# Define observation and action spaces
shape_meta:
  obs:
    image:
      shape: [3, 96, 96]  # [C, H, W]
      type: rgb
    agent_pos:
      shape: [2]          # 2D position
      type: low_dim
  action:
    shape: [2]            # 2D action (dx, dy)

# Environment for evaluation
env_runner:
  _target_: diffusion_policy.env_runner.pusht_image_runner.PushTImageRunner
  n_train: 6
  n_test: 50
  n_train_vis: 2
  n_test_vis: 4
  max_steps: 300
  n_obs_steps: ${n_obs_steps}      # Inherit from parent config
  n_action_steps: ${n_action_steps}
  fps: 10
  past_action: false

# Dataset for training
dataset:
  _target_: diffusion_policy.dataset.pusht_image_dataset.PushTImageDataset
  zarr_path: data/pusht/pusht_cchi_v7_replay.zarr
  horizon: ${horizon}
  pad_before: ${eval:'${n_obs_steps}-1'}
  pad_after: ${eval:'${n_action_steps}-1'}
  seed: ${training.seed}
  val_ratio: 0.02
  max_train_episodes: null

# Create new task config for custom environment

```

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

### DiffusionUnetImagePolicy Class Initialization

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Defines the DiffusionUnetImagePolicy class with initialization parameters for diffusion-based policy learning. Takes shape metadata, noise scheduler, observation encoder, and various horizon/step configurations as inputs. Supports FiLM conditioning through obs_as_global_cond parameter for flexible observation integration.

```Python
from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
import torch
import torch.nn as nn

class DiffusionUnetImagePolicy:
    def __init__(self,
        shape_meta: dict,
        noise_scheduler: DDPMScheduler,
        obs_encoder: nn.Module,
        horizon: int = 16,              # Prediction horizon T
        n_action_steps: int = 8,        # Action execution horizon Ta
        n_obs_steps: int = 2,           # Observation history To
        num_inference_steps: int = 100, # DDIM/DDPM sampling steps
        obs_as_global_cond: bool = True # Use FiLM conditioning
    ):
        pass
```

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

### Evaluate Diffusion Policy with PushTImageRunner

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

This Python code sets up and runs an evaluation script for a trained DiffusionUnetImagePolicy using the PushTImageRunner. It loads a policy checkpoint, configures the runner for testing, executes the policy in vectorized environments, and logs evaluation metrics (scores and videos) to Weights & Biases.

```python
# diffusion_policy/env_runner/pusht_image_runner.py
from diffusion_policy.policy.base_image_policy import BaseImagePolicy
from typing import Dict, Any
import numpy as np

class PushTImageRunner:
    def __init__(
        self,
        output_dir: str,
        n_train: int = 6,        # Training rollouts
        n_test: int = 50,        # Test rollouts
        n_train_vis: int = 2,    # Record video for N train episodes
        n_test_vis: int = 4,     # Record video for N test episodes
        max_steps: int = 300,    # Max episode length
        n_obs_steps: int = 2,
        n_action_steps: int = 8,
        fps: int = 10,
        past_action: bool = False
    ):
        """Initialize runner with vectorized environments."""
        pass

    def run(self, policy: BaseImagePolicy) -> Dict[str, Any]:
        """
        Execute policy in environments.

        Returns:
            metrics: {
                'test/mean_score': float,
                'test/sim_max_reward_0': float,
                'test/sim_video_0': wandb.Video,
                ...
            }
        """
        pass

# Complete evaluation script
import torch
from diffusion_policy.policy.diffusion_unet_image_policy import DiffusionUnetImagePolicy
from diffusion_policy.env_runner.pusht_image_runner import PushTImageRunner
import wandb

# Load checkpoint
checkpoint_path = 'data/outputs/train_0/checkpoints/latest.ckpt'
payload = torch.load(checkpoint_path, map_location='cuda:0')

# Reconstruct policy
cfg = payload['cfg']
policy = DiffusionUnetImagePolicy(cfg.policy)
policy.load_state_dict(payload['state_dicts']['model'])
policy.eval().cuda()

# Set normalizer
normalizer = payload['state_dicts']['normalizer']
policy.set_normalizer(normalizer)

# Create runner
runner = PushTImageRunner(
    output_dir='data/eval_output',
    n_train=0,      # Skip training rollouts
    n_test=50,      # 50 test episodes
    n_test_vis=10,  # Record 10 videos
    max_steps=300,
    n_obs_steps=cfg.n_obs_steps,
    n_action_steps=cfg.n_action_steps
)

# Run evaluation
with torch.no_grad():
    metrics = runner.run(policy)

print(f"Test Score: {metrics['test/mean_score']:.3f}")

# Log to wandb
wandb.init(project='diffusion_policy_eval')
wandb.log(metrics)

# Metrics structure:
# {
#   'test/mean_score': 0.92,           # Average success rate
#   'test/sim_max_reward_0': 0.95,     # Episode 0 reward
#   'test/sim_max_reward_1': 0.88,     # Episode 1 reward
#   ...
#   'test/sim_video_0': wandb.Video,   # Episode 0 video
#   'test/sim_video_1': wandb.Video,   # Episode 1 video
#   ...
# }
```

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

### Distributed Training with Ray (Python)

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Sets up and runs multi-seed parallel training for diffusion models using Ray. It defines a remote training worker function that configures and executes training for a single seed on a dedicated GPU, then launches multiple such workers across available GPUs for robust evaluation. The configuration is loaded from a YAML file and updated with seed-specific parameters.

```python
# ray_train_multirun.py
import ray
from omegaconf import OmegaConf
import os

@ray.remote(num_gpus=1)
def train_worker(config: dict, seed: int, gpu_id: int):
    """Train single seed on dedicated GPU."""
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)

    # Update config
    cfg = OmegaConf.create(config)
    cfg.training.seed = seed
    cfg.training.device = 'cuda:0'  # Will be remapped by CUDA_VISIBLE_DEVICES

    # Train
    from diffusion_policy.workspace.train_diffusion_unet_image_workspace import TrainDiffusionUnetImageWorkspace
    workspace = TrainDiffusionUnetImageWorkspace(cfg)
    workspace.run()

    return cfg.output_dir

def main():
    # Initialize Ray
    ray.init(num_gpus=4)

    # Load base config
    config = OmegaConf.load('image_pusht_diffusion_policy_cnn.yaml')
    config_dict = OmegaConf.to_container(config, resolve=True)

    # Launch 8 training runs across 4 GPUs
    seeds = [42, 43, 44, 45, 46, 47, 48, 49]
    futures = []

    for i, seed in enumerate(seeds):
        gpu_id = i % 4  # Cycle through GPUs
        future = train_worker.remote(config_dict, seed, gpu_id)
        futures.append(future)
        print(f"Launched seed {seed} on GPU {gpu_id}")

    # Wait for completion
    output_dirs = ray.get(futures)

    print("All training completed!")
    print("Output directories:", output_dirs)

    ray.shutdown()

if __name__ == "__main__":
    main()

# Aggregate metrics across seeds
# python multirun_metrics.py data/outputs/*/metrics.csv
```

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

### SequenceSampler for Temporal Sequences (Python)

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Implements a SequenceSampler for extracting overlapping temporal sequences from a replay buffer. It handles episode boundaries and padding, crucial for time-series data like robot trajectories. Dependencies include numpy.

```python
# diffusion_policy/common/sampler.py
import numpy as np

class SequenceSampler:
    def __init__(self,
        replay_buffer,
        sequence_length: int,      # T (horizon)
        pad_before: int,           # Typically n_obs_steps - 1
        pad_after: int,            # Typically n_action_steps - 1
        episode_mask: np.ndarray = None
    ):
        """
        Create sampler that handles episode boundaries.

        Args:
            pad_before: Repeat first frame this many times
            pad_after: Repeat last frame this many times
        """
        pass

    def sample_sequence(self, idx: int) -> Dict[str, np.ndarray]:
        """Sample sequence with padding."""
        pass
```

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

### Real Robot Control with Spacemouse (Python)

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Controls a real robot using a Spacemouse for teleoperation. It initializes the robot environment and Spacemouse interface, allowing users to start/stop recording and move the robot by providing delta pose commands based on Spacemouse input. Recorded data is saved to a specified directory.

```python
from diffusion_policy.real_world.real_env import RealEnv
from diffusion_policy.real_world.spacemouse import Spacemouse
import numpy as np
import time

def main(
    output_dir: str,
    robot_ip: str,
    frequency: float = 10
):
    # Create environment
    env = RealEnv(
        output_dir=output_dir,
        robot_ip=robot_ip,
        frequency=frequency,
        record_raw_video=True
    )

    # Create SpaceMouse interface
    spacemouse = Spacemouse(
        vendor_id=0x256f,
        product_id=0xc62e
    )

    print("Commands:")
    print("  C: Start recording")
    print("  S: Stop recording")
    print("  Q: Quit")

    recording = False
    dt = 1.0 / frequency

    while True:
        # Handle keyboard
        key = get_keyboard_input()  # Non-blocking

        if key == 'c' and not recording:
            print("Recording started")
            env.start_episode(time.time())
            recording = True
        elif key == 's' and recording:
            print("Recording stopped")
            env.end_episode()
            recording = False
        elif key == 'q':
            break

        # Read SpaceMouse
        motion_command = spacemouse.get_motion_state()
        # motion_command: [dx, dy, dz, drx, dry, drz] in [-1, 1]

        # Get current robot pose
        obs = env.get_obs()
        current_pose = obs['robot_eef_pose'][-1]  # Latest pose

        # Compute target pose (delta control)
        pos_delta = motion_command[:3] * 0.005  # 5mm per step at max
        rot_delta = motion_command[3:] * 0.01   # 0.01 rad per step at max

        target_pose = current_pose.copy()
        target_pose[:3] += pos_delta
        target_pose[3:] += rot_delta

        # Execute
        timestamp = np.array([time.time() + dt])
        env.exec_actions(target_pose[None], timestamp)

        time.sleep(dt)

    env.close()

# Usage
if __name__ == "__main__":
    main(
        output_dir='data/demo_pusht_real',
        robot_ip='192.168.0.204',
        frequency=10
    )

# Collected data saved to: data/demo_pusht_real/replay_buffer.zarr
```

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

### RealEnv Class for Real Robot Control (Python)

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

The RealEnv class facilitates asynchronous interaction with a real robot arm, supporting features like non-blocking observations and actions, multi-camera input, shared memory communication, and smooth motion control via RTDE. It requires robot IP, control frequency, and observation parameters during initialization.

```python
import numpy as np
from typing import Dict, Optional, List, Tuple

class RealEnv:
    def __init__(self,
        output_dir: str,
        robot_ip: str,                    # UR5 IP address
        frequency: float = 10,            # Control frequency (Hz)
        n_obs_steps: int = 2,
        obs_image_resolution: Tuple[int,int] = (640, 480),
        camera_serial_numbers: Optional[List[str]] = None,
        max_pos_speed: float = 0.25,      # m/s
        max_rot_speed: float = 0.6,       # rad/s
        tcp_offset: float = 0.13,         # Tool center point offset
        init_joints: bool = False,
        record_raw_video: bool = True,
        shm_manager = None
    ):
        """
        Asynchronous real robot environment.

        Features:
            - Non-blocking observation/action (unlike gym.Env)
            - Multi-camera RealSense support
            - Shared memory for zero-copy communication
            - RTDE interpolation controller for smooth motion
            - Automatic video recording and replay buffer saving
        """
        pass

    def get_obs(self) -> Dict[str, np.ndarray]:
        """
        Get latest observations (non-blocking).

        Returns:
            obs_dict: {
                'camera_0': (n_obs_steps, H, W, 3),
                'camera_1': (n_obs_steps, H, W, 3),
                'robot_eef_pose': (n_obs_steps, 6),  # [x,y,z,rx,ry,rz]
                'timestamp': (n_obs_steps,)          # Unix timestamps
            }
        """
        pass

    def exec_actions(
        self,
        actions: np.ndarray,      # (n_action_steps, 6) poses
        timestamps: np.ndarray    # (n_action_steps,) target times
    ):
        """
        Execute actions asynchronously (non-blocking).

        Args:
            actions: Target poses [x,y,z,rx,ry,rz] in meters/radians
            timestamps: Expected execution time for each action
        """
        pass

    def start_episode(self, start_time: float):
        """Begin recording episode."""
        pass

    def end_episode(self):
        """Stop recording and save to replay buffer."""
        pass

```

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

### LinearNormalizer for Data Normalization (Python)

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Defines a LinearNormalizer class for data normalization, supporting both 'limits' and 'gaussian' modes. It includes methods for fitting the normalizer to data, applying normalization, and reversing it. Dependencies include torch and numpy.

```python
# diffusion_policy/model/common/normalizer.py
import torch
import numpy as np

class LinearNormalizer:
    def fit(self,
        data,
        last_n_dims: int = 1,
        mode: str = 'limits',           # 'limits' or 'gaussian'
        output_max: float = 1.0,
        output_min: float = -1.0
    ):
        """
        Fit normalizer to data statistics.

        Args:
            data: torch.Tensor, np.ndarray, or dict thereof
            last_n_dims: Normalize together over last N dimensions
            mode: 'limits' scales to [output_min, output_max]
                  'gaussian' standardizes to mean=0, std=1
        """
        pass

    def normalize(self, x):
        """Apply normalization: x * scale + offset"""
        return x * self.scale + self.offset

    def unnormalize(self, x):
        """Reverse normalization: (x - offset) / scale"""
        return (x - self.offset) / self.scale
```

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

### DiffusionUnetImagePolicy Compute Loss Method

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Implements diffusion training loss computation that predicts noise at random timesteps. Takes a batch dictionary containing observations and actions, returns MSE loss between predicted and actual noise. Core training objective for the diffusion-based policy learning.

```Python
def compute_loss(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
    """
    Diffusion training loss: predict noise at random timestep.

    Args:
        batch: {'obs': {...}, 'action': (B, T, Da)}

    Returns:
        loss: MSE between predicted and actual noise
    """
    pass
```

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

### Normalize Actions and Observations with LinearNormalizer

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

This Python snippet demonstrates how to extract action and agent position data from a dataset, then normalize them using the LinearNormalizer class. It supports normalization to a specific range (e.g., [-1, 1]) or as a Gaussian distribution. The normalized data is then used to configure a diffusion policy.

```python
all_actions = []
all_agent_pos = []
for i in range(len(dataset)):
    sample = dataset[i]
    all_actions.append(sample['action'])
    all_agent_pos.append(sample['obs']['agent_pos'])

actions = np.concatenate(all_actions, axis=0)      # (N, T, Da)
agent_pos = np.concatenate(all_agent_pos, axis=0)  # (N, To, Ds)

normalizer = LinearNormalizer()

# Normalize actions to [-1, 1] independently per dimension
normalizer['action'] = LinearNormalizer()
normalizer['action'].fit(
    data=actions,
    last_n_dims=1,           # Normalize each action dim independently
    mode='limits',
    output_min=-1.0,
    output_max=1.0
)

# Normalize observations using gaussian (zero mean, unit variance)
normalizer['agent_pos'] = LinearNormalizer()
normalizer['agent_pos'].fit(
    data=agent_pos,
    last_n_dims=1,
    mode='gaussian'
)

# Use in policy
policy = DiffusionUnetImagePolicy(...)
policy.set_normalizer(normalizer)

# During training: normalization happens inside policy
loss = policy.compute_loss(batch)  # Automatically normalizes obs and actions

# During inference: automatic normalization and unnormalization
obs_dict = {'image': ..., 'agent_pos': ...}  # Raw observations
action_dict = policy.predict_action(obs_dict)  # Returns unnormalized actions
actions = action_dict['action']  # Ready to execute in environment
```

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

### Define ConditionalUnet1D Class with Timestep and Conditioning

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Initialize a 1D U-Net model for action diffusion with configurable timestep embedding, global FiLM conditioning, and local conditioning. Parameters control encoder dimensions, kernel size, and normalization. The architecture combines sinusoidal timestep encoding with MLP projection, per-channel scale/bias modulation via FiLM, and concatenative local conditioning for flexible action denoising.

```python
import torch
import torch.nn as nn
from typing import Optional, List

class ConditionalUnet1D(nn.Module):
    def __init__(self,
        input_dim: int,                          # Action dimension
        local_cond_dim: Optional[int] = None,    # For inpainting
        global_cond_dim: Optional[int] = None,   # For FiLM conditioning
        diffusion_step_embed_dim: int = 256,     # Timestep embedding
        down_dims: List[int] = [256, 512, 1024], # Encoder dimensions
        kernel_size: int = 3,
        n_groups: int = 8,                       # GroupNorm groups
        cond_predict_scale: bool = False         # FiLM with scale+bias
    ):
        """
        1D U-Net for action diffusion.

        Architecture:
            - Timestep encoding: sinusoidal + MLP
            - Global conditioning: FiLM (per-channel scale/bias)
            - Local conditioning: concatenate to input
            - Temporal convolutions with GroupNorm
            - U-Net skip connections
        """
        pass

    def forward(self,
        sample: torch.Tensor,               # (B, T, Da) noisy actions
        timestep: torch.Tensor,             # (B,) diffusion timesteps
        local_cond: Optional[torch.Tensor] = None,   # (B, T, Dc)
        global_cond: Optional[torch.Tensor] = None   # (B, Dg)
    ) -> torch.Tensor:
        """Predict noise for denoising."""
        return predicted_noise  # (B, T, Da)
```

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

### Build Diffusion Policy with Vision Encoder and ConditionalUnet1D

Source: https://context7.com/real-stanford/diffusion_policy/llms.txt

Construct a complete diffusion-based action prediction pipeline combining a ResNet18-based vision encoder with ConditionalUnet1D. The vision encoder processes multi-modal observations (RGB images and low-dimensional state) to produce global conditioning features for FiLM modulation in the noise prediction network. Includes training forward pass with DDPM noise scheduling and MSE loss computation.

```python
from diffusion_policy.model.diffusion.conditional_unet1d import ConditionalUnet1D
from diffusion_policy.model.vision.multi_image_obs_encoder import MultiImageObsEncoder
import torch

# Vision encoder (ResNet18)
obs_encoder = MultiImageObsEncoder(
    shape_meta={
        'obs': {
            'image': {'shape': [3, 96, 96], 'type': 'rgb'},
            'agent_pos': {'shape': [2], 'type': 'low_dim'}
        }
    },
    rgb_model={'_target_': 'diffusion_policy.model.vision.model_getter.get_resnet', 'name': 'resnet18'},
    crop_shape=[76, 76],
    random_crop=True,
    use_group_norm=True,  # Required for EMA
    imagenet_norm=False
)

# Get feature dimension
obs_feature_dim = obs_encoder.output_shape()[0]  # e.g., 512

# Build diffusion model
noise_pred_net = ConditionalUnet1D(
    input_dim=2,                         # 2D action space
    global_cond_dim=obs_feature_dim,     # FiLM conditioning from vision
    diffusion_step_embed_dim=128,
    down_dims=[512, 1024, 2048],
    kernel_size=5,
    n_groups=8,
    cond_predict_scale=True              # FiLM with scale and bias
)

# Training forward pass
batch_size = 64
horizon = 16

# Encode observations
obs_dict = {
    'image': torch.randn(batch_size, 2, 96, 96, 3).cuda(),
    'agent_pos': torch.randn(batch_size, 2, 2).cuda()
}
obs_features = obs_encoder(obs_dict)  # (B, obs_feature_dim)

# Add noise to actions (DDPM forward process)
actions = torch.randn(batch_size, horizon, 2).cuda()
noise = torch.randn_like(actions)
timesteps = torch.randint(0, 100, (batch_size,)).cuda()

# Predict noise
predicted_noise = noise_pred_net(
    sample=actions,           # Noisy actions
    timestep=timesteps,       # Diffusion timesteps
    global_cond=obs_features  # Conditioning
)

# Compute loss
loss = torch.nn.functional.mse_loss(predicted_noise, noise)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.