VersatIL: Imitation Learning for Any Robot Policy

🤯 The Paradox of Research Code

Have you ever found yourself wondering: "How would this Robot Policy perform if I simply swapped that ResNet18 vision encoder for an EfficientNet, or just changed the KL divergence term in the loss function for a Maximum Mean Discrepancy?"

So you clone the repo to try it out. You wrestle with a requirements.txt from 2018 that demands CUDA 9.0 and a version of PyTorch that seemingly only exists on a floppy disk in a basement. You finally get the environment running, only to discover that the loss function implementation is tightly coupled to a string variable named "dataset_v2_final_final" deep in the training loop.

Or perhaps you have wandered through State-Of-The-Art codebases, staring blankly at lines like: b = d.unsqueeze(-1).view(b, -1, h//16, w//16).permute(0, 3, 1, 2).contiguous() ...wondering what unholy things are happening to those poor tensors?

This ends with VersatIL. ⚡

VersatIL is a modular, composable framework for doing offline Imitation Learning (behavioral cloning) built with PyTorch, that decouples its main pillars: data, algorithm, network architecture and loss function into clean, reusable components.

Swap standard behavioral cloning for diffusion or flow matching, replace a CNN backbone with a ViT or a whole Vision-Language Model, or run your policy on a completely new dataset format. Compatible components can be freely swapped with config or command-line changes, no source code rewrites.

Rapid experimentation, tested and reproducible code, and true reusability across projects.

Core Principles

• 🧑‍🔬 Research-First Flexibility — Unlike frameworks that focus on reimplementing and distributing specific SOTA policies, VersatIL gives you the modular building blocks to create and benchmark your own novel architectures and algorithms on any dataset.
• 🔄 Mix & Match You are free to swap any robot policy component for easy benchmarking.
• 🧱 Modularity Each component is self-contained and reusable.
• ⚡ Modern Dependency Management – Dependencies managed with uv and pyproject.toml for modern and fast installation.
• ♻️ Don't Reinvent the Wheel We rely on industry-standard libraries:
  • Timm for vision backbones.
  • HuggingFace Transformers for Language encoders, VLMs, and tokenizers.
  • HuggingFace Diffusers for diffusion schedulers.
  • Albumentations for image augmentations.
  • torchao for post-training quantization (PT2E and quantize_() APIs).
• 💡 Invent What Matters For performance-critical components, we wrote a custom src/versatil/models/layers package in pure PyTorch. Note: These are policy-agnostic and reusable in other projects.
• 🔒 Explainability & Reproducibility – Strict interfaces, full type hints, Google-style docstrings, and runtime config validation.
• 🧪 Testing – Comprehensive unit and integration tests for every module.

Paper Reproducibility Instructions

Paper-specific instructions for reproducing reported experiments are collected in the Papers Reproducibility Instructions docs section. These notes list the datasets, local path setup, environment variables, and Hydra configs used for each paper.

🚀 Installation

VersatIL requires Python 3.13 or 3.14 and is available on PyPI:

# With uv
uv venv --python 3.14
source .venv/bin/activate
uv pip install versatil

# Or with mamba/conda
mamba create -n versatil python=3.14 pip
mamba activate versatil
pip install versatil

The default PyPI PyTorch wheel runs on both CPU-only and CUDA machines. For source installs — developing VersatIL, running the test suite, selecting the dedicated CPU-only or CUDA 13.0 PyTorch wheel sets, or enabling ExecuTorch .pte export — see the installation guide.

🏁 Quick Start

Environment Configuration

VersatIL uses a .env file to configure machine-specific paths. Copy the example and customize:

cp .env.example .env

VERSATIL_CHECKPOINT_DIR=/path/to/checkpoints      # Where model checkpoints are saved
VERSATIL_ZARR_DIR=/path/to/zarr                   # Preprocessed Zarr datasets
VERSATIL_CACHE_DIR=/path/to/cache                 # HuggingFace/torch model cache
VERSATIL_LIBERO_LEROBOT_DIR=/path/to/libero       # Raw data, one variable per dataset

The full variable list is in the installation guide.

Available Benchmarks

Ready-to-use end-to-end configs are organized by dataset under src/versatil/hydra_configs/end_to_end_training_runs/, and every simulated benchmark ships a ZMQ server wrapper for rolling out trained policies:

The LIBERO server wrapper also supports LIBERO-PRO evaluation.

Each config is self-contained, just point to your data path and run!

Training Your First Model

# Train ACT on bowel retraction dataset
python -m versatil.endpoints.train --config-name end_to_end_training_runs/bowel_retraction/act

# Train ACT on LIBERO dataset, overriding any parameter from the CLI
python -m versatil.endpoints.train --config-name end_to_end_training_runs/libero_hdf5/act \
    task.dataloader.batch_size=64 training.optimizer.lr=1e-4

CLI overrides, checkpointing, resuming, and WandB tracking are covered in the training guide.

🧩 Available Components

Encoders

Observations can be encoded with, in principle, any model hosted on timm or HF Transformers — CNNs, ViTs, language models, and VLMs alike:

• RGB: any timm vision backbone, producing spatial feature maps (ResNet, ConvNeXt, Swin, ...) or token sequences (ViT, DINOv2/v3, SigLIP, ...).
• Depth and RGB-D: single-channel depth backbones and cross-modal RGB-D encoders with geometric attention.
• Language and VLM: text embedding models (BERT-family, EmbeddingGemma, Qwen Embedding, ...) and image-text models (CLIP, SigLIP, ...) for language-conditioned policies.
• Proprioception: robot state encoders.

Features from multiple cameras and modalities can then be fused before decoding — by concatenation, MLP projection, or cross-attention. Adding a new backbone is usually a one-line Enum entry mapping to a timm or HF Transformers model name; fully custom encoders just subclass a small base interface. The complete encoder catalog is in the encoding docs.

Algorithms

The same decoder architecture can be trained with behavioral cloning, denoising diffusion, or flow matching, and any of these can be wrapped with a variational latent space learned through configurable priors and posteriors — see the algorithms docs.

Action Decoders

VersatIL reproduces the main policy architectures from the offline imitation-learning/behavioral cloning literature:

But it doesn't stop here! Beyond the literature presets, generic building blocks (bidirectional, autoregressive, and diffusion action transformers, swappable VLM backbones inside the VLA decoders, plus a Mixture-of-Experts wrapper applicable to any decoder), tuning paradigms (custom LoRA, multi-stage training, custom objective functions) and algorithm sweeping (swap algorithms from the same class, or change the diffusion/flow solver, etc.) enable you to design new architectures. You can also easily implement fully custom decoders by just subclassing ActionDecoder. The complete decoder catalog is in the decoders docs.

⏳ Workflow

📋 Before Training

⚙️ Configuration Management

VersatIL uses Hydra and OmegaConf for the management of experiment configurations.

Hydra provides:

• Hierarchical composition — Combine reusable config groups via defaults: without constant YAML editing.
• CLI overrides & sweeps — Change parameters or run hyperparameter searches directly from the command line.

OmegaConf provides:

• Config Safety: All configs are typed Dataclasses. If you pass a string where an int is expected, or forget a mandatory field, the run fails immediately at startup—not 2 hours into training.
• Smart merging — Config classes define defaults that automatically fill missing values when merging multiple files.

Together, this means you can create new experiments by overriding only the parameters that change. Forget about writing verbose YAML files for every single run you want to benchmark — an end-to-end experiment is just a list of reusable blocks:

# end_to_end_training_runs/bowel_retraction/act.yaml (excerpt)
defaults:
  - /task/dataset_schema: bowel_retraction_v2   # Raw data format
  - /task/dataloader: bowel_retraction           # Batch size, torch workers, augmentation
  - /task/action_space: deltas_cf_pos_gripper_phase
  - /task/observation_space: stereo_rgb
  - /policy/encoding_pipeline: stereo_rgb        # Encoder + fusion config
  - /policy/decoder: act_default                 # Action decoder architecture
  - /policy/algorithm: bc_with_vae_gaussian      # Learning algorithm
  - /policy/loss: regression_gripper_kl          # Loss composition

Every block is validated against a typed dataclass at startup — see the configuration guide for the full system.

🗄️ Dataset Schema (Ingestion)

Raw data formats vary wildly (Rosbags, CSVs, HDF5). We don't force you to convert your raw files manually.
Instead, VersatIL handles this with a two-stage approach:

1. DatasetSchema (how your raw data is structured) A pluggable class that maps any raw format to a standardized Zarr store.

   Schema	Class	Raw Format
   HuggingFace LeRobot	LeRobotDatasetSchemaV30	Parquet + MP4/images
   HDF5	Hdf5DatasetSchema	HDF5 archive
   CSV	CsvDatasetSchema	CSV + raw image folders
   Custom	Subclass DatasetSchema	Any

2. Zarr Store Creation Zarr provides fast, compressed, chunked storage with NumPy-like access.

   • Created automatically on first training run if missing. No separate preprocessing script needed.
   • Decouples raw storage from training-optimized layout.
   • Format-specific keys (e.g. observation.images.image, agentview_image) are remapped to standardized pipeline keys during ingestion, so training and inference never see raw format naming — see the data docs.

🏋️‍♂️ During Training

🎯 Task Definition

The TaskSpace selects what subset of the Zarr data to use, allowing multiple tasks from a single dataset without duplication:

• Observation space — Choose which observations will be given to the robot policy as state, e.g. which cameras, proprioception, depth, language instructions, etc.
• Action space — Choose which ground-truth actions will be given to your robot policy and in which format, e.g. deltas or absolute positions, gripper states, end-effector orientation, etc.
• Temporal horizons — Observation and prediction temporal windows.

Example: The same Zarr store can power a pure-vision task, a state-only task, and a vision-language task simultaneously.

🚚 Data Loading Pipeline

Uniform across all Zarr datasets:

• Fast episodic loading from Zarr
• Temporal chunking (observation windows + action sequences + masks)
• Preprocessing (normalization, augmentation, tokenization)
• Batching via PyTorch DataLoader

Actions can be precomputed (stored in Zarr) or computed on-the-fly during batching (e.g., deltas from consecutive states).

🧠 Policy Composition

A robot policy is built from four decoupled components, orchestrated by the Policy class:

1. 👁️ Encoding Pipeline: A pipeline of multi-modal encoders that extract features from raw observations plus an optional fusion module that combines the features into a unified representation.
2. 🧮 Algorithm: The learning paradigm that defines how to train the policy. This can be:

• Standard Behavioral Cloning (supervised learning of actions given observations)
• Generative approaches such as Diffusion and Flow Matching.
• Variational approaches that add a learned latent variable to any base algorithm. The latent variable can be learned through different kinds of prior-posterior schemes, which will determine the nature of the latent space.

3. 🕹️ Action Decoder: The neural architecture that decodes the features into robot actions. We provide a set of factory decoders which comprise some of the SOTA architectures of recent years in the field of offline IL, from the Action Chunking Transformer (ACT), to the Diffusion Policy, until more recent Vision-Language-Action models. More information on the available decoders can be found below.
4. 📉 Loss Module: A composable loss module that defines the objective function to optimize during training. This can be a simple regression loss (MSE) or a more complex loss that combines multiple terms (e.g. action regression + KL divergence for variational algorithms).

⚡ Training Engine

Powered by PyTorch Lightning:

• Automatic handling of training and checkpointing.
• WandB Integration: Tracks in real-time metrics and visualizations.
• Callbacks: Pluggable training-loop hooks, from EMA weights to latent-space visualizations.

🚀 Post-Training

🔌 Inference

The inference pipeline is transport-agnostic: communication with any environment server (real robot or simulation) is abstracted behind ObservationTransport and ActionTransport Python protocols. Any object satisfying these protocols works (ZMQ, HTTP, etc).

The built-in ZMQ implementation uses our two PyPI packages:

• tso-robotics-sockets: Generic ZMQ socket client/server.
• versatil-constants: Shared constants defining the observation/action message format.

Both libraries are server-agnostic, i.e. they define the message format but not the server implementation. For custom setups, you just need to implement the transport protocols, ObservationTransport and ActionTransport, and couple it with your custom server.

We also provide ZMQ server wrappers for common robot learning simulators, so VersatIL policies can be rolled out without extra glue code — each benchmark in the Available Benchmarks table links its server wrapper in the Sim Server column.

🔍 Explainability

Visual attribution for trained policies: per-camera heatmaps showing which image regions drove the predicted actions. We currently support Grad-CAM, Grad-CAM++, and Ablation-CAM, computed over every visual module of the policy (encoding-pipeline encoders as well as decoder-owned VLM vision towers), either on offline dataset samples or live during inference.

python -m versatil.endpoints.explain \
    checkpoint_path=/path/to/training/checkpoint \
    explanation_types='[gradcam]' \
    split=val \
    max_samples=16

Heatmap overlays (and optionally raw tensors) are written next to the checkpoint. See the explainability guide for sources, targeting, and output formats.

📦 Quantization and Post-Training Compression

VersatIL's quantization package is built upon PyTorch's native quantization library torchao and supports quantization-aware training as well as eager and PT2E post-training quantization, targeting the whole policy or individual modules. The post-training compression pipeline turns a trained policy checkpoint into a deployment artifact for edge or resource-constrained hardware: a PTC run can export a floating-point model, apply pruning, quantize the policy, and save either a Torch Export .pt2 artifact (x86 CPUs) or an ExecuTorch .pte artifact (ARM and x86 mobile CPUs). Background and full configuration are covered in the quantization guide and the compression guide.

🤝 Contributing

Contributions are welcome! Install from source, then verify your setup:

# Default local suite: excludes slow, integration, GPU-only, and ExecuTorch-dependent tests
pytest

# Format and lint (pre-commit hooks run this on every commit)
ruff format src/ tests/ && ruff check src/ tests/

Development setup, testing conventions, code style rules, and PR guidelines are in CONTRIBUTING.md.

🐛 Troubleshooting

Common issues (CUDA, data loading, Python 3.14 quirks) are collected in the training guide and the Known Issues page.

📖 Citation

If you use VersatIL in your research, please cite the repository:

@misc{mazza2026versatil,
    author = {Mazza, Lorenzo and Rodriguez, Ariel and Speidel, Stefanie},
    title = {VersatIL: A Modular Imitation Learning Framework for Robot Policies},
    howpublished = {\url{https://github.com/Lorenzo-Mazza/VersatIL}},
    year = {2026}
}

If you use the MoE-ACT policy, please also cite:

@article{mazza2026moe,
  title={MoE-ACT: Improving Surgical Imitation Learning Policies through Supervised Mixture-of-Experts},
  author={Mazza, Lorenzo and Rodriguez, Ariel and Younis, Rayan and Lelis, Martin and Hellig, Ortrun and Li, Chenpan and Bodenstedt, Sebastian and Wagner, Martin and Speidel, Stefanie},
  journal={arXiv preprint arXiv:2601.21971},
  year={2026}
}

If you use the LACT policy, please also cite:

@article{mazza2026understanding,
  title={Understanding Multimodal Failure in Action-Chunking Behavioral Cloning},
  author={Mazza, Lorenzo and Datres, Massimiliano and Rodriguez, Ariel and Bodenstedt, Sebastian and Kutyniok, Gitta and Speidel, Stefanie},
  journal={arXiv preprint arXiv:2605.22493},
  year={2026}
}

📄 License

VersatIL is released under the MIT License.

🙌 Contributors

Thanks to all contributors who make VersatIL possible! Want to join them? Start with the contributing guide.