Distilled Partial Chunk Critic in Offline RL

• The Distilled Partial Chunk Critic is a novel value estimation technique that decouples long-horizon TD backup from reactive, closed-loop policy execution.
• It employs optimistic expectile regression to distill maximal value predictions from a k-step chunked critic for shorter action blocks, ensuring stable learning.
• Empirical results on long-horizon tasks demonstrate that the DQC framework outperforms traditional methods by balancing efficient credit assignment with flexible policy reactivity.

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 (Li et al., 11 Dec 2025).

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: $l < k$0 Only the first $l < k$1 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 $l < k$2, the value assigned is the maximal $l < k$3 obtained by optimally completing the remainder of the chunk: $l < k$4 Direct maximization is intractable; DQC adopts a practical approximation known as "max-plus distillation," implementing implicit-max regression via expectile loss with parameter $l < k$5 to bias the estimator upwards toward the maximizing completion. This expectile regression operates as follows: $l < k$6 where $l < k$7 denotes the expectile squared-error loss, driving $l < k$8 to match the maximal achievable value obtainable from $l < k$9.

3. Loss Functions, Optimization, and Policy Extraction

The DQC architecture involves three loss functions:

• Chunked Critic TD Loss: For $k$0, standard $k$1-step TD loss is applied:

$k$2

• Partial Critic Distillation Loss: For $k$3, expectile regression towards the "teacher" $k$4 on demonstration chunks:

$k$5

• Value Head Loss: For $k$6, quantile regression (with high quantile) matches the maximal $k$7 over candidate chunks:

$k$8

The policy is not explicitly learned; instead, at test time, a set of $k$9 action chunks $a_{1:k} = (a_1,\dots,a_k)$0 are sampled from a flow-based behavior prior $a_{1:k} = (a_1,\dots,a_k)$1, and the chunk maximizing $a_{1:k} = (a_1,\dots,a_k)$2 is selected: $a_{1:k} = (a_1,\dots,a_k)$3

4. Algorithmic Workflow and Pseudocode

The full DQC algorithm for offline batch RL operates as follows:

1. For each gradient step:
   • Sample trajectory-chunks $a_{1:k} = (a_1,\dots,a_k)$4 from memory $a_{1:k} = (a_1,\dots,a_k)$5.
   • Update $a_{1:k} = (a_1,\dots,a_k)$6 parameters via $a_{1:k} = (a_1,\dots,a_k)$7-step TD loss.
   • Update $a_{1:k} = (a_1,\dots,a_k)$8 via expectile distillation from $a_{1:k} = (a_1,\dots,a_k)$9.
   • Update $k$0 via quantile regression towards $k$1.
2. At test time, for each state $k$2:
   • Draw $k$3 candidate $k$4-step chunks from $k$5.
   • Accept the chunk with maximal $k$6.

The decoupling of $k$7 (backup length) and $k$8 (policy horizon) permits the agent to plan long-term with $k$9 but act flexibly at short horizons using $Q^C(s,a_{1:k})$0.

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 $Q^C(s,a_{1:k})$1; minimum ensembling for cube tasks, mean aggregation for maze/puzzle settings.
• Chunk sizes: $Q^C(s,a_{1:k})$2 (critic), $Q^C(s,a_{1:k})$3 (policy, tuned per task).
• Loss and optimism: Expectile with $Q^C(s,a_{1:k})$4 for distillation, quantile with $Q^C(s,a_{1:k})$5 for value head.
• Sampling: Best-of-N selection with $Q^C(s,a_{1:k})$6 candidates from the flow-based prior ($Q^C(s,a_{1:k})$7 with 10 coupling steps).
• Optimization: Adam optimizer, learning rate $Q^C(s,a_{1:k})$8, batch size $Q^C(s,a_{1:k})$9 for stability, target network update $a_{1:k}$0, $a_{1:k}$1.
• 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 $a_{1:k}$2-step blocks.
• Naïve partial chunking (QC–NS), which simply executes $a_{1:k}$3 steps of the QC policy.
• Standard $a_{1:k}$4-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 $a_{1:k}$5. Varying the policy chunk length $a_{1:k}$6 showed that $a_{1:k}$7 or $a_{1:k}$8 were both effective, but large $a_{1:k}$9 recapitulated open-loop optimization hardness. The method displayed robustness to implicit loss type (expectile vs quantile) as long as mild optimism ($k$0) was maintained. Large batches ($k$1) were essential for convergence and stability on challenging tasks, with larger best-of-$k$2 values ($k$3 typical; increases to $k$4 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$5-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 (Li et al., 11 Dec 2025). 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.