Paper Rob Ibrl

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)