DifoTrain

An embodiment-agnostic framework for training Vision-Language-Action (VLA) models from the real world

---

Overview

DifoTrain is a framework for collecting demonstrations, training Vision-Language-Action (VLA) policies on them, and deploying those policies to robots — in simulation or on real hardware, on any embodiment. It learns from both teleoperation and human video, and it can either wrap an existing VLA (OpenVLA / Octo / π0 / ACT) or train a native model.

The whole collect → train → eval → deploy loop runs out of the box on CPU using a built-in, dependency-free simulated robot — no webcam, GPU, or hardware required.

Architecture

The framework hangs off four pluggable interfaces. A custom robot, data source, or model is just an implementation of one of these, registered into the framework.

Interface	Module	Implementations shipped
Robot	difotrain.embodiment	PlanarArm (sim, dependency-free); PyBulletArm, MuJoCoArm, ROS2 real-robot
DataSource	difotrain.data.sources	ScriptedTeleopSource, HumanVideoSource, SyntheticSource
Policy	difotrain.policy	Native: MLPVLAPolicy, FlowVLAPolicy (flow-matching, chunked). Wrapped: OpenVLA, ACT (LeRobot), π0 (openpi), Octo
EpisodeDataset	difotrain.data	LeRobot/RLDS-style directory dataset

Shared plumbing beneath the interfaces: a canonical ActionChunk format every policy emits, and a deterministic kinematics tokenizer that maps any (RobotSpec, observation, action chunk) into one shared vocabulary — so datasets, checkpoints and benchmark numbers stay interchangeable across embodiments.

Collect ─▶ Standardize ─▶ Train/Finetune ─▶ Evaluate ─▶ Deploy ─▶ (feedback loop)
 teleop      Episode         BC trainer        sim/real    safety     log rollouts
 human-vid   dataset         + normalize       success     layer      back to dataset

Install

pip install difotrain                 # core (numpy + torch)
pip install "difotrain[capture]"      # + MediaPipe/OpenCV for webcam human-video
pip install "difotrain[sim]"          # + PyBullet for heavier sim backends
pip install "difotrain[dev]"          # + pytest

Quickstart

The reference task: a 2-link arm must reach a named target conveyed only through language — so success measures genuine language grounding, not memorized motion.

difotrain info                                   # list registered robots/policies/sources
difotrain collect --out data/reach --episodes 200
difotrain train   --data data/reach --out runs/policy.pt --epochs 150
difotrain eval    --policy runs/policy.pt
difotrain deploy  --policy runs/policy.pt --instruction "reach to the lower left" --log-to data/rollouts
difotrain compare --data data/reach --models mlp_vla,flow_vla --out report.html

compare is the model-selection harness: it trains/evaluates every requested model on the same episodes and seeds, and emits a ranked HTML report with success rate, mean error, per-step latency and parameter count. Wrapped third-party VLAs (OpenVLA / ACT / π0 / Octo) join the comparison zero-shot when their dependencies are installed; models whose dependencies are missing become skipped rows with install instructions, never a crash. deploy --log-to writes every rollout back into a dataset, closing the collect → train → deploy → collect flywheel.

Difotrain-Bench

A versioned, reproducible benchmark suite (protocol, current leaderboard): 10 frozen tasks measuring language grounding, paraphrase robustness, observation noise, precision, out-of- distribution starts, and sim-to-sim physics transfer (PlanarArm → MuJoCo).

difotrain bench  --policy runs/policy.pt --out benchmarks/results/mine.json
difotrain export --policy runs/policy.pt --out runs/policy.onnx   # ONNX deploy
difotrain bench  --policy runs/policy.onnx                        # bench the compiled artifact
difotrain leaderboard                                             # rebuild LEADERBOARD.md

Every report JSON carries frozen seeds, the bench version, an environment fingerprint and the exact commands that reproduce the policy — submissions are PRs adding a result file (see the protocol doc).

Typical eval output after training:

success_rate     : 87.50%
mean_final_error : 0.0259
  [OK ] up           err=0.0180
  [OK ] down         err=0.0037
  ...

Python API

from difotrain.data.dataset import EpisodeDataset
from difotrain.data.sources.scripted_teleop import ScriptedTeleopSource
from difotrain.train.trainer import train_policy, TrainConfig
from difotrain.eval.evaluator import evaluate_reaching

ds = EpisodeDataset("data/reach")
ds.extend(ScriptedTeleopSource(seed=0).collect(200))     # collect demos
policy = train_policy(ds, TrainConfig(epochs=150))       # behavior cloning
print(evaluate_reaching(policy))                          # score in sim

Adding your own robot

from difotrain.core import register_robot
from difotrain.embodiment.base import Robot, RobotSpec

@register_robot("my_arm")
class MyArm(Robot):
    def __init__(self):
        self.spec = RobotSpec(name="my_arm", dof=..., observation_space=..., action_space=...)
    def reset(self, *, instruction="", seed=None): ...
    def get_observation(self): ...
    def apply_action(self, action): ...

Sim and real robots implement the same Robot API, so policies, the evaluator, and the deploy runner drive them unchanged.

Self-improvement flywheel

difotrain improve --policy runs/policy.pt --data data/reach --cycles 3

The flywheel finds the policy's failures, synthesizes corrective expert demonstrations exactly where it is weak, fine-tunes a candidate, and promotes it only if it beats the incumbent on frozen eval seeds (else automatic rollback — a cycle can never make the policy worse). difotrain sysid fits a calibrated sim twin of your robot's dynamics for training/gating. See docs/flywheel.md.

Learning from human video

HumanVideoSource estimates human pose (MediaPipe) and uses a Retargeter to map it into a robot's action space, turning cheap video into trainable demonstrations. It also runs offline from a recorded trajectory JSON:

difotrain setup     # download the MediaPipe pose model
difotrain record    # capture human motion from a webcam

Documentation

Full docs live in docs/:

• Installation
• Quickstart
• Concepts & architecture
• How it works — full architecture deep-dive
• CLI reference
• Python API
• Extending DifoTrain — add your own robot, data source, or model
• Learning from human video
• Self-improvement flywheel & sysid
• Writing an adapter — robots and model wrappers, with the conformance suite
• Kinematics tokenizer spec v1.0 — the open interchange standard
• VLA from Zero — free course
• Enterprise layer — signed bundles, fleet rollouts, federated fine-tuning
• Examples
• FAQ & troubleshooting

Roadmap

• Core interfaces, sim robot, dataset, native VLA, BC training, eval, deploy + safety
• Scripted-teleop and human-video data sources, retargeting
• PyBullet backend behind the Robot API (difotrain[sim]); MuJoCo next
• Wrapped OpenVLA / Octo / π0 / ACT policies + difotrain compare
• Difotrain-Bench frozen benchmark suite, leaderboard, ONNX export
• Self-improvement flywheel (difotrain improve), sysid twin, full safety layer
• Adapter conformance suite, tokenizer spec v1.0, model-hub packaging, course
• Enterprise layer: signed policy bundles, staged fleet rollouts, federated fine-tuning
• Image observations + transformer language encoder
• Real-hardware validation of the flywheel loop
• Isaac Lab GPU-vectorized backend (stub shipped; scene construction tracked in issues)

Development

uv sync               # or: pip install -e ".[dev,capture]"
python -m unittest discover -s tests -v

License

MIT — see LICENSE.