COIL

We're releasing the code for UniEnv, a shared abstraction layer for robotic environments and datasets, along with this codebase!

This repository accompanies the paper Correspondence-Oriented Imitation Learning: Flexible Visuomotor Control with 3D Conditioning.

• Project page: https://coil-manip.github.io/
• Paper: https://arxiv.org/abs/2512.05953

COIL is a conditional visuomotor policy learning framework built around a flexible 3D task representation: the intended motion of keypoints on scene objects, with variable spatial and temporal granularity. This code release contains the simulation-side pipeline used to generate training data, train a COIL policy, and evaluate it.

The end-to-end workflow in this repo is:

1. Collect demonstration trajectories with the scripted reference policy.
2. Train a policy on one or more saved datasets.
3. Evaluate the checkpoint, including the held-out run_sim_exp benchmark.

Setup

Requirements

• Python >=3.10
• A working PyTorch + CUDA environment for simulation/training
• A working ManiSkill / SAPIEN stack
• unienv and unienv_maniskill

The package dependencies declared in pyproject.toml include PyTorch, PyTorch3D, diffusers, accelerate, wandb, and COIL’s UniEnv wrappers. The code also imports mani_skill, so your environment must already include a compatible ManiSkill installation.

The paper’s simulation data pipeline is built on ManiSkill and UniEnv. In code, the environment factory lives in flow_manipulation/env_utils.py, and the released tasks are configured there through FlowEnvConfig.

Install the repository in editable mode:

pip install -e . --no-build-isolation

The --no-build-isolation flag is required to ensure that the pytorch3d dependency is installed correctly.

Data Collection

Demonstrations are collected by training/flow/collect_flow_traj.py. By default, it saves only successful episodes and writes a replay buffer plus per-episode context under the directory passed to --save_path.

Each dataset directory contains:

• step_data/: replay-buffer transitions used for training
• episode_context/: per-episode flow/context files
• info.json: replay-buffer metadata

Minimal collection command

Run collection from the repository root:

cd training/flow

bash collect_data.sub <task_name> <num_steps> <seed>

Or, if in a slurm system, submit the collection script with:

sbatch collect_data.sub <task_name> <num_steps> <seed>

The data collection job can sustain cancellations and will resume automatically when re-launched with the same --save_path and --seed values, when the --load flag is specified. This makes it easy to run multiple collection jobs in parallel, each with a different random seed on a pre-emptable queue.

For the COIL policy in the paper, we collected 2M steps for each task pickplace, object_sweep, and pickpour.

Training

Modify the training bash file train.sub referenced below first to point to the collected dataset directories and your slurm partition to launch jobs on. Then run it with:

cd training/flow

bash train.sub <FeatureExtractor> <seed> <batch_size>

or

sbatch train.sub <FeatureExtractor> <seed> <batch_size>

• Valid FeatureExtractor values are SmallV14FeatureExtractor, and MediumV14FeatureExtractor. In the paper, we used MediumV14FeatureExtractor.
• We used batch_size=32 in the paper.

Evaluation

There are two evaluation paths in this repository:

1. training/flow/evaluate.py evaluates the checkpoint on the selected task family with spatial correspondence generated on-the-fly through scripted policies.
2. Use this to debug your policy and get quick feedback on training progress.
3. training/flow/run_sim_exp.py runs the repo’s fixed simulation benchmark with spatial correspondence already generated with scripted policies with fixed seeds.
4. Use this to get the final numbers for the performance of your policy.

Utilities

To view the scripted reference policy in action, run:

cd scripts/
python test_env_view.py --env_config.env_name <task_name>

Citation

@misc{cao2025coil,
  title={Correspondence-Oriented Imitation Learning: Flexible Visuomotor Control with 3D Conditioning},
  author={Yunhao Cao and Zubin Bhaumik and Jessie Jia and Xingyi He and Kuan Fang},
  year={2025},
  eprint={2512.05953},
  archivePrefix={arXiv},
  primaryClass={cs.RO},
  url={https://arxiv.org/abs/2512.05953}
}