Paper Rob Ibrl

paper_rob__ibrl

Imitation Bootstrapped Reinforcement Learning (IBRL) trains a BC policy on demonstrations, then uses it alongside an RL actor during online training: at each step, a learned Q-function selects whichever action (RL or BC) has higher value for both exploration and computing TD bootstrap targets. This creates an automatic curriculum where BC dominates early and RL gradually takes over, yielding strong sample efficiency on pixel-based and state-based manipulation tasks.

Paper Info

Method Overview

1. Pretrain a BC policy on demonstration data (HDF5 format from Robomimic)
2. Initialize RL actor, dual-Q critic, and target networks
3. Warm up: collect episodes into the replay buffer using BC policy actions
4. Online RL loop: at each step, obtain both RL and BC actions, use the critic to select the better one (IBRL selection) for environment interaction
5. For TD target computation, apply the same Q-based selection on next-state actions (IBRL bootstrapping)
6. Update critic with MSE loss on Bellman targets, update actor to maximize Q, soft-update targets

Key insight: By using the BC policy to propose alternative actions for Q-based selection in both exploration and bootstrapping, IBRL leverages high-quality demonstration knowledge from the start of training without requiring explicit scheduling or regularization.

Paper-Code Mapping

Setup

Dependencies

• Python 3.9
• PyTorch 2.1.0 (CUDA 12.1)
• MuJoCo 2.1 (mujoco210 binary in ~/.mujoco/mujoco210)
• mujoco-py 2.1.2.14
• Key packages: robosuite (git), metaworld (git), pyrallis, h5py, wandb, einops, rich

Installation

# MuJoCo 2.1
wget https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz
mkdir -p ~/.mujoco && tar -xzf mujoco210-linux-x86_64.tar.gz -C ~/.mujoco/

# Clone and setup
git clone --recursive https://github.com/hengyuan-hu/ibrl.git
cd ibrl
conda create -n ibrl python=3.9 && conda activate ibrl
source set_env.sh  # sets PYTHONPATH, MUJOCO paths, OMP_NUM_THREADS=1

# Install packages
pip install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt

# Compile C++ replay buffer (pybind11)
cd common_utils && make && cd ..

Pre-trained Weights & Data

Download from Google Drive, extract to release/:

release/
├── cfgs/          # YAML configs (shipped with repo)
├── data/          # HDF5 demo datasets
│   ├── robomimic/{can,square}/processed_data96.hdf5
│   └── metaworld/{Assembly,...}/dataset.hdf5
└── model/         # Pretrained BC checkpoints
    ├── robomimic/{can,square}/model0.pt
    └── metaworld/path{Assembly,...}_*/model1.pt

Usage Scenarios

Train IBRL (pixel, Robomimic)

python train_rl.py --config_path release/cfgs/robomimic_rl/can_ibrl.yaml \
  --save_dir exp/can_ibrl --use_wb 0

Train RLPD baseline

python train_rl.py --config_path release/cfgs/robomimic_rl/can_rlpd.yaml \
  --save_dir exp/can_rlpd --use_wb 0

Train RFT baseline (two-step)

# Step 1: Pretrain RL actor with BC loss
python train_rl.py --config_path release/cfgs/robomimic_rl/can_rft.yaml \
  --pretrain_only 1 --pretrain_num_epoch 5 --load_pretrained_agent None \
  --save_dir exp/can_pretrain --use_wb 0

# Step 2: RL fine-tune with BC regularization
python train_rl.py --config_path release/cfgs/robomimic_rl/can_rft.yaml \
  --load_pretrained_agent exp/can_pretrain/model0.pt \
  --save_dir exp/can_rft --use_wb 0

Train on Meta-World

python mw_main/train_rl_mw.py \
  --config_path release/cfgs/metaworld/ibrl_basic.yaml \
  --bc_policy assembly --save_dir exp/mw_ibrl --use_wb 0

Key Config Flags

Code Integration Guide

Minimal Imports

import sys
sys.path.append("/path/to/ibrl")

from rl.q_agent import QAgent, QAgentConfig
from bc.bc_policy import BcPolicy, BcPolicyConfig
from networks.encoder import VitEncoder, VitEncoderConfig

Model Instantiation & Inference

import torch
from rl.q_agent import QAgent, QAgentConfig

# Create QAgent for pixel-based task
cfg = QAgentConfig()
cfg.act_method = "ibrl"           # or "rl", "ibrl_soft"
cfg.enc_type = "vit"
cfg.use_prop = 1
cfg.vit.embed_style = "embed2"
cfg.vit.depth = 1
cfg.actor.dropout = 0.5
cfg.critic.spatial_emb = 1024

agent = QAgent(
    use_state=False,
    obs_shape=(3, 96, 96),         # (C*obs_stack, H, W)
    prop_shape=(27,),              # prop_dim * prop_stack (9 * 3)
    action_dim=7,                  # OSC_POSE: 6 EE + 1 gripper
    rl_camera="robot0_eye_in_hand",
    cfg=cfg,
)

# Load trained weights
agent.load_state_dict(torch.load("model0.pt"))
agent.eval()

# Inference: obs is dict with camera images and proprioception
# Images: uint8 [C, H, W], will be normalized internally (/255 - 0.5)
# Prop: float32 [prop_dim]
obs = {
    "robot0_eye_in_hand": torch.randint(0, 255, (3, 96, 96)).cuda(),
    "prop": torch.randn(27).cuda(),
}
with torch.no_grad():
    action = agent.act(obs, eval_mode=True)  # returns [7] tensor on CPU

Loading a Full Trained Model

import train_rl

agent, eval_env, eval_env_params = train_rl.load_model("path/to/model0.pt", "cuda")
# agent: QAgent with BC policy attached (if bc_policy was set in config)
# eval_env: PixelRobosuite ready for evaluation

Data Format

Integration Notes

• The repo must be on PYTHONPATH due to absolute imports (import common_utils, from rl.q_agent import ...)
• common_utils has a C++ pybind11 module (rela) that must be compiled with make before use
• The QAgent owns separate optimizers internally; for custom training loops, access agent.encoder_opt, agent.critic_opt, agent.actor_opt
• BC policies are stored in agent.bc_policies list and must be added via agent.add_bc_policy()
• Image observations must have pixel values > 5 (raw uint8 range); normalization is done inside the encoder
• source set_env.sh is critical: it sets PYTHONPATH, MuJoCo paths, and OMP_NUM_THREADS=1 (needed for parallel eval)

Core Architecture

IBRL Training Loop
==================

Demonstrations (HDF5)
  |
  v
BC Policy Training (train_bc.py)
  |
  +-- BcPolicy: MultiViewEncoder(per-cam ResNet + compress + cat) -> MLP -> tanh
  |   or
  +-- StateBcPolicy: MLP(dropout=0.5) -> tanh
  |
  v
RL Training (train_rl.py)
  |
  +-- QAgent (rl/q_agent.py)
  |     |
  |     +-- Encoder: VitEncoder / ResNetEncoder / DrQEncoder
  |     |     - Input: [B, C, 96, 96] uint8 -> /255 - 0.5
  |     |     - Output: [B, num_patch, patch_dim]
  |     |
  |     +-- Actor: feat -> compress(Linear+LN+ReLU or SpatialEmb) -> MLP -> TruncatedNormal
  |     |     - dropout=0.5 on RL actor helps prevent overfitting
  |     |
  |     +-- Critic: dual-Q, feat+action -> _QNet or SpatialEmbQNet -> scalar
  |     |     - Target: min(Q1_target, Q2_target)
  |     |
  |     +-- IBRL Selection (act + bootstrap):
  |           a_rl = Actor(s), a_bc = BCPolicy(s)
  |           a* = argmax_{a in {a_rl, a_bc}} Q_target(s, a)  [hard]
  |           or: sample from softmax(Q * beta)                [soft]
  |
  +-- Replay Buffer (C++ pybind11)
  |     - N-step returns (default n=3)
  |     - Separate BC replay for demo mixing (RLPD: mix_rl_rate=0.5)
  |
  +-- Data Augmentation: RandomShiftsAug(pad=4)

State-based variant:
  Actor = FcActor(MLP, dropout=0.5)
  Critic = MultiFcQ(10 Q-nets via _MultiLinear, RED-Q: sample k=2 for min)

Repo Structure

Three Training Paradigms

Tips & Gotchas

• C++ compilation required: cd common_utils && make is mandatory; requires cmake, gcc, and pybind11 (included as submodule). If you see GLIBCXX_3.4.30 not found, symlink system libstdc++ into conda env
• source set_env.sh: must run once per shell; sets PYTHONPATH=$PWD, MuJoCo paths, and OMP_NUM_THREADS=1 (critical for parallel eval performance)
• MuJoCo 2.1 specifically: the repo uses mujoco-py which requires the old MuJoCo 2.1 binary (not the newer mujoco pip package alone)
• GPU memory: pixel-based ViT training uses ~16GB VRAM; ResNet/DrQ encoders use less
• Warm-up phase: the first 50 episodes use BC policy actions to fill the replay buffer; ensure bc_policy path is valid
• Actor dropout: 0.5 dropout on the RL actor is important for IBRL; it prevents overfitting to demonstration-heavy early data. During critic target computation, actor_target.training is asserted True (dropout active)
• assert False at end: both train_rl.py and train_bc.py deliberately crash at the end (assert False) to signal completion; this is intentional
• Image normalization: encoders expect raw pixel values (>5), normalization /255 - 0.5 happens inside the encoder forward pass
• Evaluation parallelism: --mp_eval 1 spawns 10 processes; requires OMP_NUM_THREADS=1 to avoid thread contention
• Config system: pyrallis supports YAML file + CLI overrides; nested fields use dots: --q_agent.actor.dropout 0.3
• Action space: Robomimic uses OSC_POSE (6D end-effector delta + 1 gripper), actions clamped to [-1, 1] via Tanh output
• HDF5 data format: demos stored as data/demo_N/obs/{camera}_image, data/demo_N/actions, data/demo_N/rewards
• Meta-World task names: use short names (assembly, boxclose, coffeepush, stickpull) which map to BC_POLICIES and BC_DATASETS dicts in mw_main/train_rl_mw.py
• update_freq: 2: RL update happens every 2 env steps; critic updates num_critic_update times, actor updates only on the last critic iteration (RED-Q style)