Paper Rl Qc

paper_rl__qc

Q-Chunking (QC) applies action chunking to temporal-difference RL methods, enabling agents to operate in a temporally extended action space with unbiased n-step backups. This improves exploration and online sample efficiency in offline-to-online settings, particularly on long-horizon, sparse-reward manipulation tasks.

Paper Info

Method Overview

1. Action chunking: The actor outputs horizon_length actions concatenated into a single vector. The critic evaluates Q(s, a_0:H) where a_0:H is the full chunk. During execution, the entire chunk is committed open-loop.
2. N-step returns: With chunk length H, the Bellman target uses the cumulative discounted reward over H steps and bootstraps from the state H steps later, with discount raised to the power H: r_0 + gamma*r_1 + ... + gamma^(H-1)*r_{H-1} + gamma^H * Q(s_H, a_H).
3. Offline-to-online: Offline pre-training uses main.py (ACFQL agent). Online fine-tuning continues with main.py (same script, runs offline then online phases) or uses main_online.py (ACRLPD agent, pure online with replay buffer mixing).
4. Two agent variants: ACFQL (flow Q-learning with flow-matching actor) and ACRLPD (SAC-based with TanhNormal actor).

Key insight: Action chunking applied to TD-based RL provides both structured exploration (by committing to multi-step action sequences from offline data) and correct value learning (via unbiased n-step returns), bridging the gap between behavior cloning's temporal coherence and RL's optimality.

Paper-Code Mapping

Setup

Dependencies

• Python 3.10+
• JAX 0.6.0+ (with CUDA 12 support)
• Flax 0.10.5+
• Key packages: jax, flax, optax, ml_collections, distrax, tensorflow-probability, wandb, ogbench, mujoco
• For RoboMimic tasks: robomimic, h5py
• For D4RL tasks: d4rl

Installation

git clone https://github.com/ColinQiyangLi/qc.git
cd qc
pip install -r requirements.txt

Datasets

OGBench (auto-downloaded):

• Standard datasets download automatically via ogbench.make_env_and_datasets()
• Large datasets (e.g., cube-quadruple-play-100m-v0): manually download and use --ogbench_dataset_dir=<path>

# Download cube-quadruple 100M dataset
wget -r -np -nH --cut-dirs=2 -A "*.npz" https://rail.eecs.berkeley.edu/datasets/ogbench/cube-quadruple-play-100m-v0/

RoboMimic (manual download):

• Place datasets at ~/.robomimic/<task>/<type>/low_dim_v15.hdf5
• Download from https://robomimic.github.io/docs/datasets/robomimic_v0.1.html (Multi-Human MH links)
• Supported env name format: <task>-<type>-low_dim (e.g., lift-mh-low_dim, can-mh-low_dim, square-mh-low_dim)

Usage Scenarios

Offline-to-Online RL (QC with best-of-N actor)

MUJOCO_GL=egl python main.py \
    --agent=agents/acfql.py \
    --agent.actor_type=best-of-n \
    --agent.actor_num_samples=32 \
    --env_name=cube-triple-play-singletask-task2-v0 \
    --horizon_length=5 \
    --offline_steps=1000000 \
    --online_steps=1000000 \
    --seed=0

Offline-to-Online RL (QC-FQL with distill-DDPG actor)

MUJOCO_GL=egl python main.py \
    --agent=agents/acfql.py \
    --agent.actor_type=distill-ddpg \
    --agent.alpha=100 \
    --env_name=cube-triple-play-singletask-task2-v0 \
    --horizon_length=5 \
    --seed=0

Pure Online RL (QC-RLPD with BC regularization)

MUJOCO_GL=egl python main_online.py \
    --agent=agents/acrlpd.py \
    --env_name=cube-triple-play-singletask-task2-v0 \
    --horizon_length=5 \
    --agent.bc_alpha=0.01 \
    --seed=0

Sparse Reward Tasks

MUJOCO_GL=egl python main.py \
    --agent.actor_type=best-of-n \
    --agent.actor_num_samples=32 \
    --env_name=scene-play-singletask-task0-v0 \
    --sparse=True \
    --horizon_length=5 \
    --seed=0

Baselines (no action chunking)

# FQL (horizon_length=1, no chunking, no n-step)
MUJOCO_GL=egl python main.py --agent.alpha=100 --env_name=cube-triple-play-singletask-task2-v0 --horizon_length=1

# BFN-n (n-step returns without chunking, actor sees single action)
MUJOCO_GL=egl python main.py --agent.actor_type=best-of-n --agent.actor_num_samples=4 \
    --env_name=cube-triple-play-singletask-task2-v0 --horizon_length=5 --agent.action_chunking=False

# RLPD (no chunking, standard SAC online)
MUJOCO_GL=egl python main_online.py --env_name=cube-triple-play-singletask-task2-v0 --horizon_length=1

Key Config Flags

Code Integration Guide

Minimal Imports

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

from agents.acfql import ACFQLAgent, get_config as get_acfql_config
from agents.acrlpd import ACRLPDAgent, get_config as get_acrlpd_config
from utils.datasets import Dataset, ReplayBuffer
from evaluation import evaluate

Agent Instantiation

import jax.numpy as jnp

# Get default config and customize
config = get_acfql_config()
config["horizon_length"] = 5
config["actor_type"] = "best-of-n"
config["actor_num_samples"] = 32

# Create agent from example data
# ex_observations: shape (obs_dim,) e.g. (23,) for RoboMimic lift
# ex_actions: shape (action_dim,) e.g. (7,)
agent = ACFQLAgent.create(
    seed=0,
    ex_observations=ex_observations,
    ex_actions=ex_actions,
    config=config,
)

Training Loop (single step)

import jax

# Offline update (single batch)
batch = dataset.sample_sequence(
    batch_size=256,
    sequence_length=5,     # horizon_length
    discount=0.99,
)
agent, info = agent.update(batch)

# Online update (batched UTD)
utd_ratio = 4
batch = replay_buffer.sample_sequence(
    batch_size=256 * utd_ratio,
    sequence_length=5,
    discount=0.99,
)
batch = jax.tree.map(lambda x: x.reshape((utd_ratio, 256) + x.shape[1:]), batch)
agent, info = agent.batch_update(batch)

Action Sampling and Execution

import jax
import numpy as np

rng = jax.random.PRNGKey(0)
action = agent.sample_actions(observations=obs, rng=rng)
# action shape: (horizon_length * action_dim,) when action_chunking=True
# Reshape and execute sequentially:
action_chunk = np.array(action).reshape(-1, action_dim)
for a in action_chunk:
    next_obs, reward, terminated, truncated, info = env.step(a)

Data Format

Integration Notes

• The codebase uses JAX/Flax, not PyTorch. Agents are Flax PyTree nodes (immutable).
• Agent updates return a new agent: agent, info = agent.update(batch). You must reassign.
• ModuleDict in utils/flax_utils.py bundles all networks (actor, critic, target_critic) into a single TrainState with a shared optimizer.
• TrainState.select(name) returns a partial call that routes to a specific sub-module.
• The batch_update method uses jax.lax.scan over the UTD dimension for efficient multi-update.
• Config is passed as ml_collections.ConfigDict via absl flags. Agent config files double as the config source (get_config() at module bottom).
• MUJOCO_GL=egl is required for headless GPU rendering (cluster environments).
• ACFQL uses utils/networks.py (MLP with GELU, ensemblize). ACRLPD uses rlpd_networks/ and agents/model.py (MLP with ReLU, nn.vmap Ensemble). They are not interchangeable.
• Save/restore via pickle: save_agent(agent, dir, step) writes params_<step>.pkl, restore_agent_with_file(agent, path) loads it.

Core Architecture

                    ACFQL (Flow Q-Learning)                    ACRLPD (SAC-based)
                    ~~~~~~~~~~~~~~~~~~~~~~                     ~~~~~~~~~~~~~~~~~~

Observation ----+---> ActorVectorField (BC flow)              Observation ---> TanhNormal Actor
                |         |                                                        |
                |    [Euler integration, flow_steps=10]                       sample actions
                |         |                                                        |
                |    flow actions (H*action_dim)                             actions (H*action_dim)
                |         |                                                        |
                |    +----+---- (distill-ddpg)                               +-----+
                |    |         ActorVectorField (one-step)                    |
                |    |              |                                         |
                |    +---- OR ---- (best-of-n)                               |
                |    |         sample N, pick max Q                          |
                |    |                                                       |
                +---> Value (critic ensemble, num_qs=2)        Ensemble(StateActionValue, num_qs=10)
                |         |                                         |
                |    Q(s, a_0:H)                              Q(s, a_0:H)
                |         |                                         |
                +---> Value (target critic, EMA tau=0.005)    target critic (EMA tau=0.005)
                          |                                         |
                     target_q = sum(gamma^i * r_i)            target_q = sum(gamma^i * r_i)
                              + gamma^H * mask * Q_target           + gamma^H * mask * Q_target
                                                                         |
                                                                    + alpha * entropy
                                                                    + bc_alpha * BC_loss

Dataset.sample_sequence(batch_size, sequence_length=H, discount=gamma)
    -> {observations, actions[B,H,A], rewards[B,H], masks[B,H], valid[B,H], next_observations[B,H,...]}

Repo Structure

Tips & Gotchas

• action_chunking=True vs False: When True, the actor outputs H*action_dim actions and the critic takes the concatenated chunk as input. When False, the actor outputs a single action but still uses n-step returns (the "n-step return baseline" in the paper).
• horizon_length=1: Reduces to standard single-step RL (no chunking, no n-step return).
• Two different network stacks: ACFQL uses utils/networks.py (MLP with GELU activation, ensemblize via nn.vmap). ACRLPD uses rlpd_networks/ and agents/model.py (MLP with ReLU, separate Ensemble class). They are not interchangeable.
• Reward adjustment: RoboMimic rewards are shifted by -1.0 (from [0,1] to [-1,0]). D4RL AntMaze rewards are also shifted by -1.0. Sparse mode further transforms to (r != 0) * -1.0.
• Evaluation action queue: During eval, a new action chunk is only sampled when the queue is empty. Actions are popped FIFO. The queue is cleared on episode reset.
• WandB required by default: Both scripts initialize WandB at startup. Set WANDB_MODE=disabled to run without logging.
• Memory for large datasets: Use --ogbench_dataset_dir with --dataset_replace_interval=1000 to cycle through dataset shards (avoids loading entire 100M dataset into memory).
• EGL rendering: Set MUJOCO_GL=egl on headless machines. The code also auto-sets EGL_DEVICE_ID and MUJOCO_EGL_DEVICE_ID from CUDA_VISIBLE_DEVICES.
• JAX JIT compilation: First training step will be slow due to JIT compilation. Subsequent steps are fast.
• Valid mask in sequence sampling: The valid field marks whether each position in the H-step window is before any terminal. The critic loss is weighted by valid[..., -1] to ignore invalid target bootstraps.
• main.py vs main_online.py: main.py has both offline and online phases (replay buffer initialized from offline data). main_online.py is online-only with an empty replay buffer, mixing 50/50 dataset batches + replay batches at each step.
• batch_update vs update: update processes a single batch. batch_update takes a batch with an extra leading dimension (UTD ratio) and scans over it with jax.lax.scan, returning averaged info.
• RoboMimic env naming: Use <task>-<dataset_type>-low_dim format, e.g., lift-mh-low_dim. The is_robomimic_env() check requires low_dim in the name. Max episode lengths: lift/can=300, square=400, transport=800, tool_hang=1000.
• Checkpoint format: Saved as params_<step>.pkl via pickle with flax.serialization.to_state_dict().
• Online action queue reset: When an episode terminates during online training, the action queue is cleared and a new chunk is sampled at the next step.