Dense-Jump Flow Matching Policy

This repository provides a minimal, reference implementation of the Dense-Jump Flow Matching Policy. The method introduced is designed to be simple and effective, making it easy to implement and extend—perfect for "vibe coding" tools or integrating on top of your own flow matching policy codebase.

The main purpose of this repository is to share additional experiments and ablation studies that were not included in the original paper due to length restrictions.

Installation

First, clone the repository and navigate into the project directory:

git clone https://github.com/DenseJumpFM/DenseJump_FlowMatching.git

pip install -r requirements.txt

Datasets

This project primarily uses datasets from Gymnasium and Minari.

• Gymnasium: https://gymnasium.farama.org/ - A standard API for reinforcement learning environments.
• Minari: https://minari.farama.org/ - A library for hosting and managing offline reinforcement learning datasets.

Repository Structure & Usage

1. Training (train_flow.py)

This is the main entry point for training the flow matching policy. It handles data loading (supporting hdf5 and pkl), normalization, and the training loop. This will store every 100 epoch as checkpoints.

Usage Example:

python train_flow.py 
    --data_path “your data path".pkl \
    --num_epochs 5000 \
    --obs_hidden_dims 128 128 \
    --policy_hidden_dims 256 256 \
    --lr 1e-4 \
    --beta 0.2 0.2\ # change to --beta 1.0 1.0 for uniform sampling
    --batch_size 128 \
    --seed 0 \
    --num_train 8000

2. Inference (inference.py)

Used to evaluate a trained policy in its environment. It will run all checkpoints in the model path directory and will can export results to a CSV file. It loads task configurations from tasks_paths.json.

Usage Example:

python inference.py 
    --data_path "your data path".pkl \
    --policy_path "your policy path" #can be a singe model path or a directory of checkpoints \
    --num_envs 100 \
    --num_episode 1 \
    --seed 0 \
    --euler_steps "[0 0.5]" # this is the time steps you want to use when integrate

5. Analysis Tools (knn/)

This directory contains scripts for analyzing the learned vector fields, primarily used for the ablation studies mentioned in the paper.

a. Build KNN Index (knn/knn_build_index.py)

Builds a k-Nearest Neighbors index on the dataset actions. This is required for the validation script to efficiently query the training distribution.

Usage:

python knn/knn_build_index.py \
  --data_path  "your data path"\
  --output knn/knn_index_pendulum.joblib \
  --num_train   # number of training data you used

b. Validate Velocity Field (knn/validate_velocity_field.py)

Validates the velocity field generated by the flow matching policy relative to the training data. It computes metrics like cosine similarity between the predicted velocity and the direction towards the true target (or nearest neighbor), helping to verify if the flow is correctly guiding samples to the data manifold.

Usage:

python knn/validate_velocity_field.py \
  --data_path "your data path" \
  --policy_path "your policy path" # use single model path\
  --knn_index knn/knn_index_pendulum.joblib \
  --num_train  # number of training data you used\
  --euler_steps "[0, 0.5]" \
  --plot_cos
 

Citation

If you find this work useful, please cite our paper:

@inproceedings{chen2025densejump,
  title={Dense-Jump Flow Matching with Non-Uniform Time Scheduling for Robotic
Policies: Mitigating Multi-Step Inference Degradation},
  author={Chen, Zidong and Guo, Zihao and Wang, Peng and Egbe, ThankGod Itua and Lyu, Yan and Qian, Chenghao},
  booktitle={2026 IEEE International Conference on Robotics and Automation (ICRA)},
  year={2026},
  url={https://arxiv.org/abs/2509.13574}
}