ReinFlow: Fine-tuning Flow Matching Policies

ReinFlow supports Weights & Biases (WandB) for experiment tracking.

Key Features:

• Offline Mode: Use wandb.offline_mode=True to log locally without internet. Sync later using wandb sync <directory>.
• Disable Logging: Set wandb=null to skip WandB creation and save memory.
• Data Recovery: If wandb=null was used, training records are still saved as .pkl files. You can recover them using util/pkl2wandb.py.

# Sync offline WandB runs in a batch
for dir in ./wandb_offline/wandb/offline-run-20250510*; do wandb sync "${dir}"; done

ReinFlow interacts with environments using the OpenAI Gym format. To add a new environment, ensure it is compatible with the following patterns:

• Vectorization: Environments are initialized as vectorized environments via make_async (found in env/gym_utils/__init__.py).
• History & Multi-step: The multi_step wrapper (env/gym_utils/wrapper/multi_step.py) is used to handle history observations and multi-environment-step action execution.
• Normalization: Use environment-specific wrappers (like robomimic_lowdim or furniture) to handle observation and action normalization. You can implement a new wrapper to support your specific environment's requirements.

IsaacGym is a required dependency. You can download the installer from the NVIDIA website or use the provided AWS S3 link. After downloading, unzip the package and install it in editable mode using pip.

Note: It is recommended to use the --no-cache-dir and --force-reinstall flags during installation to avoid common issues. Do not upgrade pip if prompted, as the current version is required for codebase compatibility.

# Download the package
cd ~
wget https://iai-robust-rearrangement.s3.us-east-2.amazonaws.com/IsaacGym_Preview_4_Package.tar.gz

# Unzip the package
tar -xzf IsaacGym_Preview_4_Package.tar.gz

# Install the package
cd isaacgym
pip install -e python --no-cache-dir --force-reinstall

For Franka Kitchen and Robomimic datasets, the pre-training commands handle downloading and normalization automatically.

• Franka Kitchen: Uses D4RL datasets (complete, mixed, or partial observations) via DPPO's approach.
• Robomimic: Uses DPPO's simplified pixel datasets. Note that these are smaller and simpler than official Robomimic data; using official data may require larger models and more training time but can improve success rates.

To use MuJoCo with ReinFlow, you must install the MuJoCo 2.1.0 binaries and configure your environment variables. This involves downloading the tarball, extracting it to ~/.mujoco, and updating your ~/.bashrc to include the MuJoCo binary directory in your LD_LIBRARY_PATH and PATH.

cd ~
wget https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz
mkdir $HOME/.mujoco
tar -xvf mujoco210-linux-x86_64.tar.gz -C ~/.mujoco/
echo -e 'export LD_LIBRARY_PATH=$HOME/.mujoco/mujoco210/bin' >> ~/.bashrc
echo -e 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidia' >> ~/.bashrc
echo -e 'export PATH="$LD_LIBRARY_PATH:$PATH"' >> ~/.bashrc

DPPO's version of the datasets is older and less recommended. It downloads automatically when use_d4rl_dataset=False. Use this only if you need to match exact results from DPPO research. If you use both D4RL and DPPO datasets, store them in separate directories and manually update train_dataset_path and normalization_path in your execution commands.

python script/run.py --config-dir=cfg/gym/pretrain/walker2d-medium-v2 --config-name=pre_reflow_mlp device=cuda:0 sim_device=cuda:0

Furniture-Bench is a required dependency. Install it by cloning the repository (via SSH or HTTPS) and installing it in editable mode using pip install -e ..

# Using SSH
git clone git@github.com:ankile/furniture-bench.git

# OR using HTTPS
git clone https://github.com/ankile/furniture-bench.git

# Install the package
cd furniture-bench
pip install -e .

To recreate the plots from the ReinFlow paper, you must first obtain the necessary data files. Download all files from the visualize branch of the ReinFlow repository and unzip them into the <REINFLOW_DIR>/visualize/ directory.

mkdir -p <REINFLOW_DIR>/visualize/
# Unzip files from 'visualize' branch into this directory