Implicit Behavioral Cloning (IBC)

Comparing Implicit Behavioral Cloning against MSE regression on multimodal control tasks

Energy-based imitation learning with InfoNCE training and DFO + Langevin inference

---

Overview

This repository re-implements Implicit Behavioral Cloning (IBC) from the paper Implicit Behavioral Cloning (Florence et al., CoRL 2021) and compares it against a classic MSE regression baseline on two environments:

Environment	Description	Scripts
Particle 2D	Toy multimodal task — a particle must visit two goals in random order	scripts/train_ibc.py, scripts/train_mse.py, scripts/evaluate.py, scripts/visualize.py
Block push	Real robot data from the IBC paper — push a block to a target location	scripts/train_block_push.py, scripts/evaluate_block_push.py, scripts/visualize_block_push.py, scripts/dagger_block_push.py

💡 Key takeaway: on multimodal data, MSE regression averages over the modes and produces meaningless actions, while IBC learns an energy landscape and can pick a single valid mode.

---

Key Idea: IBC vs MSE

Aspect	MSE Regression	IBC (Implicit Behavioral Cloning)
Model	ExplicitMLP — maps obs → action directly	EBMMLP — energy function E(obs, action)
Training	Mean-squared error between predicted and expert action	InfoNCE loss: expert action must have lower energy than sampled negatives
Inference	Single forward pass	Optimize the lowest-energy action via DFO + Langevin dynamics (DFOLangevinSampler in inference.py)
Multimodal data	Averages over modes → invalid action	Selects one mode correctly

⚠️ Note: MSE is a unimodal regressor by construction — when the same observation maps to several valid actions, it collapses to their mean. IBC instead models the full conditional distribution implicitly through the energy function.

---

Pipeline

  Expert demos ──► Preprocess (TFRecord → NPZ) ──► Train IBC / MSE ──► Evaluate ──► Visualize

• Particle 2D: demos are generated on the fly by an oracle policy (envs/particle_env.py).
• Block push: raw TFRecord data in data/block_push_states_location/ is converted once to data/block_push_processed.npz by preprocess.py.

---

Installation

pip install -r requirements.txt

Package	Purpose
numpy	Data handling & sampling
torch	Model training & inference
matplotlib	Static plots & video animation
tfrecord	Reading raw block push TFRecords

---

Usage

1. Preprocess block push data (once)

python preprocess.py

Note: This reads the raw TFRecords in data/block_push_states_location/ and writes the normalized dataset to data/block_push_processed.npz.

2. Train particle 2D

python scripts/train_ibc.py
python scripts/train_mse.py

3. Train block push (IBC, MSE, or both)

python scripts/train_block_push.py --model both

--model accepts ibc, mse, or both (default: both). Checkpoints are saved to checkpoints/ (block_push_ibc.pt, block_push_mse.pt).

4. Evaluate

python scripts/evaluate.py
python scripts/evaluate_block_push.py

scripts/evaluate.py reports success rates for Oracle vs MSE vs IBC on the multimodal particle task, and demonstrates DFO convergence on a single initial state.

5. Visualize

python scripts/visualize.py
python scripts/visualize_block_push.py

scripts/visualize_block_push.py runs closed-loop rollouts of the expert and IBC policies in envs/block_push_sim.py, saving comparison plots to plots/ and an animation to videos/.

---

Demo

Video — expert vs IBC (block push)

💡 Note: if the video does not render in your viewer, open videos/block_push_compare.gif directly.

---

Project Structure

ibc/
├── models/                    # Network architectures
│   ├── ebm_mlp.py             # EBMMLP — energy function for IBC
│   └── explicit_mlp.py        # ExplicitMLP — direct obs→action for MSE
├── envs/                      # Environments & data loaders
│   ├── particle_env.py        # Particle 2D env + oracle (multimodal demos)
│   ├── block_push_loader.py   # TFRecord loader for block push data
│   └── block_push_sim.py      # Block push closed-loop simulator
├── training.py                # train_ibc, train_mse, info_nce_loss, val_accuracy
├── inference.py               # DFOSampler, LangevinSampler, DFOLangevinSampler, load_model
├── preprocess.py              # TFRecord → data/block_push_processed.npz
├── paths.py                   # Central path definitions
├── scripts/                   # Entry points
│   ├── train_ibc.py           # Train IBC on particle 2D
│   ├── train_mse.py           # Train MSE on particle 2D
│   ├── train_block_push.py    # Train block push (--model ibc|mse|both)
│   ├── evaluate.py            # Particle: Oracle vs MSE vs IBC
│   ├── evaluate_block_push.py # Block push evaluation
│   ├── visualize.py           # Particle visualization
│   ├── visualize_block_push.py# Block push plots + video
│   └── dagger_block_push.py   # DAgger-style training for block push
├── data/                      # Raw TFRecord + block_push_processed.npz
├── checkpoints/               # Trained models
│   ├── ibc_ebm.pt             # Particle IBC
│   ├── mse_mlp.pt             # Particle MSE
│   ├── block_push_ibc.pt      # Block push IBC
│   ├── block_push_mse.pt      # Block push MSE
│   └── block_push_ibc_dagger.pt # Block push IBC (DAgger)
├── plots/                     # PNG comparison plots
├── videos/                    # MP4 animations
└── requirements.txt

---

References

• Implicit Behavioral Cloning — Florence, Lynch, Zeng, Ramirez, Wahid, Downs, Wong, Lee, Mordatch, Tompson. CoRL 2021.
• Deep Energy-Based Policies for Reinforcement Learning — Haarnoja, Tang, Abbeel, Levine. ICML 2017.
• Official IBC repository (google-research/ibc)

---