Going with the Flow: Koopman Behavioral Models as Implicit Planners for Visuo-Motor Dexterity

This repository contains the official implementation of the paper:

Going with the Flow: Koopman Behavioral Models as Implicit Planners for Visuo-Motor Dexterity
📄 https://arxiv.org/abs/2602.07413

---

Overview

We propose Koopman United Behavioral Models (KUBMs), a unified framework for learning visuo-motor manipulation policies as coupled dynamical systems over robot actions and visual observations.

Unlike conventional reactive policies that predict only the next action, our method:

• learns a latent linear dynamical system (via Koopman Operator theory)
• enables full-horizon rollout of behavior trajectories
• performs reactive planning by need via monitoring the sudden observation changes

This results in:

• temporally coherent actions
• built-in mechanisms for self-monitoring and replanning

---

Illustration of Koopman Unified Behavioral Models (KUBM)

---

Repository Structure

.
├── data/                # datasets and preprocessing
│   ├── DexArt/          # DexArt dataset
│   │   ├── Training_set/       # training set
│   │   └── Test_set/        # test set
│   │
│   ├── Adroit/          # Adroit dataset
│   │   ├── Training_set/       # training set
│   │   └── Test_set/        # test set
│
├── training/
│   ├── KUBM/            # our method KUBM
│   ├── Diffusion/       # Diffusion
│   ├── ACT/             # ACT
│   └── Flow_training/   # Training Flow autoencoder
│
├── evaluation/          # trained models and simulation evaluation
│   ├── dexart/          # SAPIEN runners for bucket, laptop, toilet
│   ├── adroit/          # MuJoCo runners for door, hammer, pen, relocate
│   ├── checkpoints/    # task autoencoders, KUBM models, configs, normalization
│   ├── configs/        # per-task inference YAML files
│   └── third_party/    # perception and custom simulator implementations
│
├── assets/              # figures, videos, teaser
│
├── environment.yaml     # conda environment for reproducibility
│
└── README.md

Installation

1. Clone the repository

git clone https://github.com/GT-STAR-Lab/K-UBM.git
cd K-UBM
conda env create -f environment.yaml
conda activate kbm

2. Download datasets

The datasets used in this project are not included in this repository due to size constraints.

Please download the datasets from the following link:

• DexArt & Adroit datasets:
  https://drive.google.com/file/d/17ZE-xOFVBMg8Jz0qWvLhVwH6HAp9V4Wg/view?usp=sharing

Setup

After downloading, unzip the file and organize the data under the following structure:

data/
├── DexArt/
│   ├── Training_set/
│   └── Test_set/
├── Adroit/
│   ├── Training_set/
│   └── Test_set/

Training

Flow Autoencoder Training

The flow autoencoder is used to learn compact latent representations of visual flow features.
For each task, the precomputed flow features are stored under the corresponding images/ folder in the dataset directory.

Files

• autoencoder_training.py
  Trains the flow autoencoder on flow points extracted from the training set using SAM3 and CoTracker.

The example training command is:

python3 training/Flow_training/autoencoder_training.py \
  --root_dir ./data/DexArt/Training_set/bucket/images/ \
  --config_path ./training/KUBM/config_autoencoder.yaml \
  --save_dir ./data/DexArt/Training_set/bucket/images/

• get_features.py
  Extracts flow features using the trained autoencoder.

The example command is:

python3 training/Flow_training/get_features.py \
  --root_dir ./data/DexArt/Training_set/bucket/images/ \
  --config_path ./training/KUBM/config_autoencoder.yaml \
  --model_dir ./data/DexArt/Training_set/bucket/images/<model_path>  *(replace with the path to the trained model)*

Train K-UBM

After preparing the action trajectories and flow features, train K-UBM with:

python3 training/KUBM/visual_kodex_training_NN_lifting.py \
  --act_dir ./data/DexArt/Training_set/bucket/actions/ \
  --flow_dir ./data/DexArt/Training_set/bucket/images/ \
  --frame_skip 1 \
  --config_path ./training/KUBM/config_NN_KODex.yaml

---

Note: Trained models are saved to
./data/<Dataset>/Training_set/<Task>/images/networks/<seed>/
(default seed in config_NN_KODex.yaml: 42).

Evaluate K-UBM

Simulation evaluation is separated by environment:

Folder	Simulator	Tasks and selected seeds
evaluation/dexart/	SAPIEN	bucket 15, laptop 25, toilet 25
evaluation/adroit/	MuJoCo / MJRL	door 15, hammer 20, pen 10, relocate 25

Both folders include the matching flow-autoencoder and KUBM model definitions. Their simulator/control loops are task-family specific. Checkpoints, normalization arrays, inference configs, and shared SAM3/CoTracker implementations are under evaluation/.

After preparing the appropriate simulator, assets, dataset, and perception dependencies, example commands from the repository root are:

# DexArt: bucket
python evaluation/dexart/Visual_KODex_inference.py --inference_setting evaluation/configs/KODex_inference_bucket.yaml --seed 15 --video

# Adroit / DAPG: door
python evaluation/adroit/Visual_KODex_inference.py --inference_setting evaluation/configs/KODex_inference_door.yaml --seed 15 --visualize True --video

Append --test_set for the test split. See the evaluation guide for all tasks and the Adroit guide for MuJoCo setup, external mesh assets, goal-feature data, and rendering requirements.

Both task families include reference reactivity scripts. Their task-specific experiments and current limitations are documented in the evaluation guides; they are not generic live replanning demos. Checkpoint/CPU compatibility checks passed, but full simulator/perception rollouts have not been validated for this packaged layout.