Sequence Compression Speeds Up Credit Assignment in Reinforcement Learning

TL;DR

This work tackles the temporal credit-assignment problem in reinforcement learning with long-delayed, stochastic outcomes. It introduces Chunked-TD, a model-based approach that uses forward-model predictions to dynamically chunk trajectories and compute adaptive $\lambda$-returns, enabling online, faster bootstrapping while remaining robust to model inaccuracies. The method combines history compression with an online eligibility-trace formulation, and it encompasses variants for state-value, action-value, and factored-reward settings. Across synthetic environments, Chunked-TD demonstrates faster credit assignment and reduced variance compared with standard TD$\lambda$, including strong performance in factored-reward tasks, highlighting its practical potential for efficient learning in discrete-action domains.

Abstract

Temporal credit assignment in reinforcement learning is challenging due to delayed and stochastic outcomes. Monte Carlo targets can bridge long delays between action and consequence but lead to high-variance targets due to stochasticity. Temporal difference (TD) learning uses bootstrapping to overcome variance but introduces a bias that can only be corrected through many iterations. TD($λ$) provides a mechanism to navigate this bias-variance tradeoff smoothly. Appropriately selecting $λ$ can significantly improve performance. Here, we propose Chunked-TD, which uses predicted probabilities of transitions from a model for computing $λ$-return targets. Unlike other model-based solutions to credit assignment, Chunked-TD is less vulnerable to model inaccuracies. Our approach is motivated by the principle of history compression and 'chunks' trajectories for conventional TD learning. Chunking with learned world models compresses near-deterministic regions of the environment-policy interaction to speed up credit assignment while still bootstrapping when necessary. We propose algorithms that can be implemented online and show that they solve some problems much faster than conventional TD($λ$).

Figures (8)

• Figure 1: Two-level Chunker-RNNs can shorten credit assignment paths. Only the tokens associated with incorrect predictions (denoted with red) at the lower level are passed on to the higher level RNN (with the start token). The hidden state of the RNN at level $l$ and index $i$ is denoted as $h^l_i$.
• Figure 2: Sample-based chunking (see Section \ref{['sec:method:subsec:nstep']}) for state value estimation. All next states (and rewards) are deterministic apart from the outcome at $x_k$, where $P_{\pi}(x_{k+1} | x_k)$ is low. The states in red are a sample from our algorithm, i.e., chunking based on model probabilities that provide targets for bootstrapping. The chunked target for $s_2$ (with $\gamma=1$, denoted as $G_2^C$) is shown as an example.
• Figure 3: Chain-and-Split environment. States are denoted by large circles and actions by lines leading to the small black circles. Colored states have a reward upon entering that state.
• Figure 4: Results from the Chain-and-Split environment. The action-value gap ($\Delta Q$) between optimal action $a_1$ and the maximum action value among the remaining actions. The true value of $\Delta Q = 0.01$. Shading indicates 95% bootstrapped confidence intervals over 10 independent trials.
• Figure 5: Results from the Accumulated-Charge environment. The average number of regretful choices over 10 runs for different algorithms with different learning rates for the accumulated-charge experiment. The dashed blue line represents the best possible expected regretful choices under the exploration scheme.

Theorems & Definitions (2)

• Proposition 1
• proof