---
title: Distilled Partial Chunk Critic in Offline RL
url: https://www.emergentmind.com/topics/distilled-partial-chunk-critic
type: topic
---

# Distilled Partial Chunk Critic in Offline RL

The Distilled Partial Chunk Critic is a value estimation approach in offline reinforcement learning that forms the core of the Decoupled Q-Chunking (DQC) framework. It enables multi-step value propagation by learning a critic over long open-loop action chunks, then distilling a partial-chunk critic that enables closed-loop, reactive policy extraction over much shorter action blocks. This design allows agents to benefit from efficient long-horizon temporal-difference (TD) backup without incurring the optimization and expressivity challenges associated with open-loop policies for long action sequences [2512.10926].

## 1. Formal Definitions and Notation

The chunked critic framework defines two separate chunk lengths: the critic chunk length $k$ and the policy chunk length $l < k$. An action chunk of length $k$ is written $a_{1:k} = (a_1,\dots,a_k)$. The $k$-step chunked critic, $Q^C(s,a_{1:k})$, estimates the expected discounted return from executing $a_{1:k}$ open-loop, then continuing optimally in future $k$-chunks:
\[
Q^{C}(s_{t},a_{t:t+k-1}) \approx \mathbb{E}\left[\sum_{i=0}^{k-1}\gamma^i r(s_{t+i},a_{t+i}) + \gamma^k \max_{a'_{t+k:t+2k-1}} Q^{C}(s_{t+k},a'_{t+k:t+2k-1})\right].
\]
The Distilled Partial Chunk Critic, denoted $Q^P(s,a_{1:l})$, approximates the maximal value achievable by optimally completing a given partial chunk:
\[
Q^P(s,a_{1:l}) \approx \max_{a_{l+1:k}} Q^C(s,[a_{1:l},a_{l+1:k}]).
\]
Only the first $l$ actions are produced and executed per decision step (closed-loop policy), and the process iterates.

## 2. Distilled Partial-Chunk Backup and Optimistic Regression

The key insight is to perform an "optimistic" backup: for any fixed partial action chunk $a_{1:l}$, the value assigned is the maximal $Q^C$ obtained by optimally completing the remainder of the chunk:
\[
Q^P(s,a_{1:l}) \leftarrow \max_{a_{l+1:k}} Q^C(s,[a_{1:l},a_{l+1:k}]).
\]
Direct maximization is intractable; DQC adopts a practical approximation known as "max-plus distillation," implementing implicit-max regression via expectile loss with parameter $\kappa > 0.5$ to bias the estimator upwards toward the maximizing completion. This expectile regression operates as follows:
\[
L_{\text{distill}}(\psi) = \mathbb{E}_{(s,a_{1:k})\sim D}\left[f^{\kappa}_{\text{imp}}(Q^C_{\bar{\phi}}(s,a_{1:k}) - Q^P_{\psi}(s,a_{1:l}))\right],
\]
where $f^{\kappa}_{\text{imp}}$ denotes the expectile squared-error loss, driving $Q^P$ to match the maximal achievable value obtainable from $Q^C$.

## 3. Loss Functions, Optimization, and Policy Extraction

The DQC architecture involves three loss functions:
- **Chunked Critic TD Loss:** For $Q^C$, standard $k$-step TD loss is applied:
  \[
  L_{\rm QC}(\phi) = \mathbb{E}\left[\left(Q_{\phi}(s,a_{1:k}) - (R_{0:k-1}+\gamma^{k} V(s_k))\right)^2\right].
  \]
- **Partial Critic Distillation Loss:** For $Q^P$, expectile regression towards the "teacher" $Q^C$ on demonstration chunks:
  \[
  L_{\rm distill}(\psi) = \mathbb{E}\left[f^{\kappa_d}_{\text{expectile}}(Q^C_{\phi}(s,a_{1:k}) - Q^P_{\psi}(s,a_{1:l}))\right].
  \]
- **Value Head Loss:** For $V_\xi$, quantile regression (with high quantile) matches the maximal $Q^P$ over candidate chunks:
  \[
  L_{V}(\xi) = \mathbb{E}\left[f^{\kappa_b}_{\text{quantile}}(\bar{Q}^P_{\psi}(s,a_{1:l}) - V_{\xi}(s))\right].
  \]

The policy is not explicitly learned; instead, at test time, a set of $N$ action chunks $\{a_{1:l}^i\}_{i=1}^N$ are sampled from a flow-based behavior prior $\pi_\beta(\cdot|s)$, and the chunk maximizing $Q^P$ is selected:
\[
a^{*} = \arg\max_{i} Q^P_{\psi}(s,a_{1:l}^i).
\]

## 4. Algorithmic Workflow and Pseudocode

The full DQC algorithm for offline batch RL operates as follows:
1. For each gradient step:
   - Sample trajectory-chunks $(s_t...s_{t+k}, a_t...a_{t+k-1}, r_t...r_{t+k-1})$ from memory $D$.
   - Update $Q^C$ parameters via $k$-step TD loss.
   - Update $Q^P$ via expectile distillation from $Q^C$.
   - Update $V_\xi$ via quantile regression towards $Q^P$.
2. At test time, for each state $s$:
   - Draw $N$ candidate $l$-step chunks from $\pi_\beta$.
   - Accept the chunk with maximal $Q^P$.

The decoupling of $k$ (backup length) and $l$ (policy horizon) permits the agent to plan long-term with $Q^C$ but act flexibly at short horizons using $Q^P$.

## 5. Implementation Details and Architectural Choices

Key instantiations in DQC are as follows:
- **Network architecture:** 4-layer MLP, 1024 units/layer, with ReLU nonlinearities.
- **Critic ensembles:** Size $K=2$; minimum ensembling for cube tasks, mean aggregation for maze/puzzle settings.
- **Chunk sizes:** $k=25$ (critic), $l \in \{1,5\}$ (policy, tuned per task).
- **Loss and optimism:** Expectile with $\kappa_d \approx 0.8$ for distillation, quantile with $\kappa_b \approx 0.97$ for value head.
- **Sampling:** Best-of-N selection with $N=32$ candidates from the flow-based prior ($\pi_\beta$ with 10 coupling steps).
- **Optimization:** Adam optimizer, learning rate $3 \times 10^{-4}$, batch size $4096$ for stability, target network update $\lambda = 5 \times 10^{-3}$, $\gamma=0.999$.
- **Engineering practices:** Large batches ensure expectile regression stability and mild upward bias in the critic; the action-flow prior regularizes the chunk sampling distribution.

## 6. Empirical Results and Comparative Analysis

DQC and the Distilled Partial Chunk Critic were evaluated on six OGBench long-horizon tasks: cube-triple, cube-quadruple, cube-octuple, humanoidmaze-giant, puzzle-4×5, and puzzle-4×6. Performance comparisons demonstrated that DQC outperformed:
- Open-loop Q-chunking (QC) using full $k$-step blocks.
- Naïve partial chunking (QC–NS), which simply executes $l$ steps of the QC policy.
- Standard $n$-step return TD (NS) and 1-step TD (OS).
- Prior state-of-art offline RL approaches including SHARSA, IQL, HIQL, FBC, and HFBC.

Ablation studies indicated that excluding the distilled critic (QC-NS variant) sharply reduced performance, particularly with $l>1$. Varying the policy chunk length $l$ showed that $l=1$ or $l=5$ were both effective, but large $l$ recapitulated open-loop optimization hardness. The method displayed robustness to implicit loss type (expectile vs quantile) as long as mild optimism ($\kappa > 0.5$) was maintained. Large batches ($\geq 4096$) were essential for convergence and stability on challenging tasks, with larger best-of-$N$ values ($N=32$ typical; increases to $128$ provided little further gain).

## 7. Significance and Practical Implications

The Distilled Partial Chunk Critic enables a principled decoupling between the benefits of long-horizon, multi-step TD backup (efficient credit assignment), and the closed-loop, reactive nature of deep RL policy execution. By learning a partial chunk critic through optimistic max regression from a $k$-step chunked critic, and extracting policies via best-of-N sampling, DQC sidesteps the memorization and optimization bottlenecks inherent in open-loop chunk policies for long blocks. A plausible implication is improved scalability to more complex, sparse-reward, or high-horizon tasks without incurring prohibitive policy search or representational complexity [2512.10926]. This framework represents a significant development for offline RL, particularly in contexts where long-term planning is essential, yet policy flexibility and reactivity cannot be compromised.

Source: https://www.emergentmind.com/topics/distilled-partial-chunk-critic