amlink/so101-diffusion-policy-resnet18-baseline

SO101 Diffusion Policy baseline — dual ResNet18 + Mish U-Net

这是 amlink 自行训练的 SO101 双相机仿真策略，供后续蒸馏、概率动作头和策略加速实验复用。主权重为通过完整测试的 step_010000.pt EMA，不是官方发布的预训练 SO101 策略，也不是 SenseCraft/LeRobot 标准导入包。

This is a locally trained simulation policy using pinned official Diffusion Policy components. It takes two RGB cameras and measured joint states, and produces a 16-step chunk of absolute joint targets. No language input or simulator object state is supplied to the network.

What is included

• step_010000.pt: original, unmodified selected EMA checkpoint, 1,316,760,254 bytes. Includes configuration, TRAIN normalization statistics and all model weights; no Adam optimizer.
• checkpoint_config.json, normalization_stats.json: readable exports from that exact checkpoint.
• code/model.py: original model wrapper, unchanged from the accepted experiment.
• code/vendor/: minimal unmodified official component source, pinned SHA manifest and upstream MIT license.
• load_baseline.py, infer_npz.py: portable loading / offline tensor inference helpers.
• evidence/metrics.json, a representative successful video, and validation/portable_loader.json.

Checkpoint SHA-256:

ade3b5100fa63cded3cfadf836a39d5fa5f7ddf435fc68d7630c7aab3898b156

Architecture and training

Item	Configuration
Parameters	329,158,918, FP32
Visual encoder	Two independent ResNet18 backbones, GroupNorm, ImageNet input normalization
Initial visual weights	Random initialization (weights=None), trained on this simulation dataset
Visual input	Previous/current frame × front/gripper RGB × 96×96
Condition	2060 features: two frames of both RGB encodings plus 6-joint proprioception
Action model	Official ConditionalUnet1D, Mish, channels [512,1024,2048], kernel 5, FiLM scaling/bias
Training diffusion	DDPM, 100 noise timesteps, cosine beta schedule, epsilon prediction
Training	256 demonstrations, 24 validation episodes; seed 42; 100,000 updates; batch 64
Optimizer	AdamW 1e-4, betas [.95,.999], decay 1e-6; 500 warmup steps and cosine schedule
Selected weights	EMA at completed update 10,000; all ten periodic checkpoints tied at 24/24, so the predeclared earliest-checkpoint rule applied
Main evaluation	DDPM100; action horizon H=16, execute K=8 at 10 Hz

The full training run completed 100,000 updates. Selecting the 10,000-step checkpoint does not mean training ended at 10,000. The final 100,000-step checkpoint is not the default evaluated baseline and is intentionally not substituted here.

Measured results and limits

Evaluation date: 2026-10-27. Task: SO101 grasps a 24 mm / 15 g rigid cube and releases it onto a green support pad, with a narrow randomized initial-position range.

Sampler, same selected weights	Validation	New test	Historical paired test	RTX 5090 full inference median
DDPM100 (baseline)	24/24	100/100	100/100	1328.99 ms
DDIM16	24/24	Not evaluated in this release	Not evaluated in this release	211.85 ms
DDIM32	24/24	Not evaluated in this release	Not evaluated in this release	419.71 ms
DDIM50	24/24	Not evaluated in this release	Not evaluated in this release	653.59 ms

New-test seeds were 15000–15099, historical paired seeds 13000–13099. The new-test Wilson 95% success interval is approximately 96.30%–100%. This is one training seed, not evidence of universal success or real-world generalization.

Physical success requires sustained bilateral grasp, sustained lift above 45 mm, placement within 22 mm, open gripper with no jaw contact, positive support force, low linear/angular speed, and ten stable control steps. These predicates were independently recomputed from saved telemetry; a success reward alone was not used as sufficient evidence.

Physics pauses during inference. DDPM100 exceeds the 800 ms replanning budget (K8 at 10 Hz), so its 100/100 score is not a real-time deployment claim. DDIM16 is the faster comparison variant, but its 100-episode test was not part of the accepted release. Timing includes CPU CHW inputs → GPU transfer → encoder → denoising → decode/clamp/finite check → synchronous CPU output; it excludes camera capture, rendering, physics, serial communication and RNG reseeding.

Representative successful test video, seed 15000

Download

This repository is public. No account token is required to download it.

pip install huggingface_hub
hf download amlink/so101-diffusion-policy-resnet18-baseline --local-dir so101-dp-baseline

For reproducible experiments, add --revision <commit_sha> using the revision shown in Files and versions or the upload handoff. Verify provenance.json and artifact_manifest.json.

Load and use as a teacher

Use Python 3.11 or newer and an isolated environment. Install a compatible PyTorch/torchvision pair for your platform, then:

cd so101-dp-baseline
python -m pip install -r requirements.txt

The original Linux GPU evaluation used Python 3.12, torch 2.11.0+cu128, torchvision 0.26.0+cu128, diffusers 0.35.2 and einops 0.8.2. The separately recorded portable CPU loading test identifies its actual versions; it is a compatibility check, not a repeated task evaluation. requirements.txt does not overwrite your installed PyTorch/CUDA build.

import torch
from load_baseline import load_baseline

device = "cuda:0"  # use "cpu" for loading/offline checks
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
model, checkpoint = load_baseline(device)

# Supply real task observations; shapes and units below are mandatory.
# images: uint8 [B, 2, 2, 3, 96, 96]
# q_env: float32 [B, 2, 6], environment-normalized measured joint positions
generator = torch.Generator(device=device).manual_seed(401)
with torch.inference_mode():
    actions_env = model.predict(
        images.to(device), q_env.to(device),
        sampler="ddpm", inference_steps=100, generator=generator,
    )  # [B, 16, 6], environment-normalized absolute targets

# Use actions_env[:, :8] in the original simulator's bounded controller,
# then update the two-frame observation history and replan.

For action-head research:

with torch.no_grad():
    context = model.encode(images.to(device), q_env.to(device))  # [B, 2060]
    noise = torch.randn((len(context), 16, 6), device=device)
    teacher_target = model.head_predict(
        context, noise, sampler="ddim", inference_steps=16
    )  # TRAIN-minmax action coordinates, before deployment clipping
    actions_env = model.decode(teacher_target.clamp(-1, 1)).clamp(-1, 1)

DDIM uses eta=0. DDPM consumes additional Gaussian draws during reverse steps: fixing only the initial noise is not enough to reproduce DDPM; seed/pass its generator too. Always label the sampler used to generate distillation targets and report it separately from DDPM100 task quality.

For a supplied NPZ observation batch:

python infer_npz.py --input observations.npz --output actions.npz --device cuda:0 --sampler ddim --steps 16

This helper performs tensor inference only; it has no serial or robot connection.

Exact observation and action contract

• Image axes are [batch, previous/current, front/gripper, RGB, height, width]. Images must be RGB uint8, not BGR or already ImageNet-normalized floats. The camera layout is near_top_wider_wrist_v3; changing the camera pose/domain is not an equivalent input.
• Joint order: shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper.
• Input q and output actions are not degrees, raw servo registers or radians. The simulator first uses q_env = 2*(q_rad-low)/(high-low)-1, where low/high are the intersection of the original joint and actuator ranges. The network applies an additional TRAIN min/max transformation internally, stored in normalization_stats.json.
• model.predict includes TRAIN-range clipping → inverse TRAIN scaling → environment-range clipping. head_predict returns values before these deployment transformations.
• Action index 0 belongs to the current observation. Execute indices 0–7, not a history-offset slice. The original simulator also limits each commanded joint change to 0.08 rad per 100 ms.
• Two frames must be the most recent actual observations, not two identical frames on every replan. Duplicating the first observation is only the initial-history convention.
• This checkpoint is specific to the simulated scene and action semantics. It is not a drop-in real-SO101 calibration/controller or SenseCraft model ZIP.

Reproduction and provenance

The full training/evaluation project is rayggxia-workspace/robot-policy-sim, commit ad554d14fdec957b8fdce2d4c865b1d3404e8b17, experiment experiments/so101_official_dp. Project repository access may be required for the full simulator/dataset workflow; the model loading code needed here is included in this public model repository. Full training data, optimizer state, robot meshes and simulator dependencies are not part of this upload.

Official source: real-stanford/diffusion_policy, revision 5ba07ac6661db573af695b419a7947ecb704690f. The vendored subset retains its upstream MIT license at code/vendor/LICENSE and per-file hashes. This was trained locally from scratch, not redistributed third-party pretrained weights. Upstream component licenses do not automatically grant a blanket license to unrelated experiment assets.