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.

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/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"

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.

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
• 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
• Image observations + transformer language encoder
• Real-hardware drivers and online DAgger feedback loop

Development

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

License

MIT — see LICENSE.