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/          # evaluation environments require SAPIEN and MuJoCo to be installed.
│
├── 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

python3 visual_kodex_test_NN_lifting.py \
  --act_dir ./data/DexArt/Test_set/bucket/actions/ \
  --flow_dir ./data/DexArt/Test_set/bucket/images/only_cur_features/ \
  --frame_skip 1 \
  --model_dir ./data/DexArt/Training_set/bucket/images/networks/42