---
title: Return-Conditioned Reinforcement Learning
url: https://www.emergentmind.com/topics/return-conditioned-reinforcement-learning-rcrl
type: topic
---

# Return-Conditioned Reinforcement Learning

Return-Conditioned Reinforcement Learning (RCRL) is a family of methods that cast policy learning as conditional supervised learning or conditional sequence modeling: instead of learning only $\pi(a\mid s)$, they fit a policy conditioned on a desired return, such as $\pi(a\mid s,g)$ or $\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$, and then query that policy at test time with a high-return target [2206.01079][1912.13465]. In the offline setting, the core training signal is an offline dataset of trajectories labeled by cumulative return or return-to-go, and the central promise is that suboptimal trajectories can still provide valid supervision for matching their observed return, while the central difficulty is that conditioning on return does not by itself remove coverage, stochasticity, or trajectory-stitching constraints [2206.01079].

## 1. Formalism and policy class

A standard RCRL setup assumes an offline dataset $\mathcal D$ of trajectories
\[
\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),
\]
with return
\[
g(\tau)=\sum_{t=1}^H r_t.
\]
Using any state summary $s_t$, RCRL fits a conditional action model $\pi(a\mid s,g)$ by minimizing the empirical negative log-likelihood
\[
\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),
\]
and then defines a test-time policy by choosing a conditioning function $f(s)$ and executing
\[
\pi_f(a\mid s)=\pi(a\mid s,f(s)).
\]
In the infinite-data optimum, if $\pi(a\mid s,g)$ can represent the true behavior-conditional distribution $P_\beta(a\mid s,g)$, the induced policy is
\[
\pi_f^{\mathrm{RCSL}}(a\mid s)=P_\beta(a\mid s,f(s))
=\beta(a\mid s)\frac{P_\beta(g=f(s)\mid s,a)}{P_\beta(g=f(s)\mid s)}.
\]
This form makes explicit that RCRL re-weights behavior actions according to their likelihood of achieving the queried return [2206.01079].

Closely related formulations differ mainly in conditioning granularity. Sequence-modeling variants define
\[
R_t=\sum_{t'=t}^{T} r_{t'}
\]
and train
\[
\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)
\]
by maximum likelihood over offline trajectories [2312.13716]. Other formulations write a time-indexed policy
\[
\pi:S\times[H]\times \mathbb R\to \Delta(A),
\]
with loss
\[
\hat L(\pi)=-\sum_{\tau\in D}\sum_{h=1}^H \log \pi(a_h\mid s_h,h,g(\tau,h)),
\]
so that test-time execution uses $\pi_f(a\mid s,h)=\pi(a\mid s,h,f(s,h))$ [2506.08463]. In the reward-conditioned formulation of behavior policy $\beta$, the reward-to-go random variable is
\[
Z^\beta(s,a)=\sum_{t\ge 0}\gamma^t r_t \;\Bigl|\; s_0=s,\;a_0=a,\;a_{t>0}\sim\beta,\;s_{t+1}\sim P,
\]
and one trains $\bar\beta_\theta(a\mid s,R)$ on tuples $(s_i,a_i,R_i)$, executing a high target RTG at test time and decrementing it as rewards are observed [2305.11340].

## 2. Canonical formulations and architectures

An early formulation appears in “Reward-Conditioned Policies,” which models a single policy
\[
\pi_\theta(a\mid s,Z)
\]
conditioned on a scalar “value target” $Z$, either the discounted return-to-go or the estimated advantage $A(s,a)$ [1912.13465]. The paper derives the method from a KL-constrained policy-search objective, obtaining a non-parametric optimum
\[
p_{\pi^*}(\tau,Z)\propto p_\mu(\tau,Z)\exp(Z/\beta),
\]
with the interpretation that, conditioned on achieving return $Z$, the optimal policy reproduces the data-collection trajectories that achieved $Z$, while the target distribution over $Z$ is shifted toward higher values by $\exp(Z/\beta)$ [1912.13465]. The resulting algorithmic family includes RCP-R and RCP-A, replay-buffer relabeling, optional sample weighting by $\exp(Z/\beta)$, and multiplicative conditioning rather than simple concatenation to prevent the network from ignoring $Z$ [1912.13465].

Transformer instantiations place RCRL inside autoregressive sequence modeling. Decision Transformer-style models tokenize a trajectory into triples $(g_t,s_t,a_t)$, embed return, state, and action with learned embeddings, and train a Transformer to predict the next action from a masked history of returns, states, and actions [2410.23450]. In another formulation, the encoder ingests $(s_t,R_t,a_{t-1})$ with positional encodings and causal masking, and predicts $\hat a_t\sim \pi_\theta(\cdot\mid \tau_{0:t-1},s_t,R_t)$ [2312.13716]. More generally, reward conditioning “enables use of generic architectures (MLPs, Transformers, diffusion models) but carries no explicit inductive bias about how different $R$ values relate” [2305.11340].

Control Transformer extends the sequence-modeling view to long-horizon robot navigation. It represents each navigation trajectory as
\[
(\hat R_0,s_0,g_0,a_0,r_0,\dots,\hat R_T,s_T,g_T,a_T,r_T),
\]
where $g_t$ is the goal offset and $\hat R_t$ is the return-to-go, then learns an autoregressive policy
\[
\pi_\theta(a_t\mid \hat R_{t-k:t},s_{t-k:t},a_{t-k:t-1},g_{t-k:t})
\]
and a goal-conditioned value function $V_\phi(s\mid g)$ used at test time to approximate the unknown optimal initial return [2211.06407]. In that system, planning trajectories are produced by a PRM-guided low-level controller, relabeled with final-goal offsets, and used as offline data for return-conditioned sequence modeling [2211.06407].

## 3. Guarantees, assumptions, and structural limitations

The sharpest theoretical analysis in the supplied literature studies when RCRL recovers a near-optimal policy in offline RL. Under three assumptions—return coverage, near determinism, and consistency of the conditioning function—the suboptimality of the ideal RCSL policy is bounded by
\[
\mathbb E_{s_1}[f(s_1)]-J(\pi_f^{\mathrm{RCSL}})\le \epsilon H^2(1/\alpha_f+2),
\]
and there exists an $f$ such that
\[
J(\pi^*)-J(\pi_f^{\mathrm{RCSL}})\le \epsilon H^2(1/\alpha_f+3).
\]
If $\epsilon=0$ and $f(s_1)=V^*(s_1)$ with $\alpha_f>0$, then $J(\pi_f^{\mathrm{RCSL}})=J(\pi^*)$ [2206.01079]. These results show exact recovery in fully deterministic settings, but only under assumptions stronger than those needed by dynamic-programming-based offline RL [2206.01079].

The same analysis gives counterexamples establishing necessity. In the “One-arm” example, $\alpha_f=\epsilon$, but the $\pi_f$ policy cannot produce return $1$, so suboptimality $1-\epsilon$ matches $\epsilon/\alpha_f=1$. In “Max-arm,” conditioning on $g=1$ leads RCSL to always pick the red arm, incurring regret approximately $1-2\epsilon$, which shows bias from stochasticity. In “Irremediable bias,” regardless of $f$, RCSL picks arms $50$–$50$, producing irreducible suboptimality [2206.01079]. A further negative result is sample-complexity: in deterministic MDPs the required number of samples can be exponential in $H$ if $\alpha_f\sim O(e^{-cH})$ [2206.01079].

A distinct theoretical line emphasizes why RCSL is attractive relative to off-policy dynamic programming when function approximation is imperfect. RCSL “never applies $B$; it performs a one-shot supervised fit to the observed $(s,g,a)$ triples,” and therefore avoids the Bellman completeness assumption that DP methods require for stable convergence [2310.19308]. The “Width Separation” theorem states that there exists a family of deterministic MDPs such that any two-layer ReLU network approximating both $Q^*$ and closed under the Bellman operator requires hidden-layer width $\Omega(u)$, whereas a two-layer ReLU network implementing an RCSL policy $\pi(a\mid s,g)$ needs width $O(1)$ independent of $u$ [2310.19308]. At the same time, purely data-driven RCSL lacks stitching on suboptimal datasets: one impossibility theorem gives a constant gap for Markovian RCSL, and another states that a Decision-Transformer-style RCSL trained on a deterministic dataset with no expert rollouts cannot achieve the optimal return under a mild generalization assumption [2310.19308].

## 4. Bayesian and critic-guided corrections

Several extensions modify vanilla RCRL to address two recurring pathologies: poor generalization for high return-to-go and inconsistencies induced by stochastic environments. “Bayesian Reparameterization of Reward-Conditioned Reinforcement Learning with Energy-based Models” identifies two critical challenges: improving generalization on high RTG inputs and avoiding out-of-distribution RTG queries during testing time [2305.11340]. The paper attributes a core obstacle to “RTG Independence,” meaning that vanilla RCRL feeds $R$ as an input token and treats $p(a\mid s,R_i)$ and $p(a\mid s,R_j)$ as independent prediction tasks [2305.11340]. Its proposed decomposition
\[
\beta(a\mid s,Z^\beta(s,a)=R)\propto \beta(a\mid s)\beta(R\mid s,a)
\]
leads to an energy-based model with
\[
E_\theta(s,a;R)=-\log \bar\beta_\theta(a\mid s)-\log \bar\beta_\theta(R\mid s,a),
\]
trained by a base conditional-likelihood loss plus an auxiliary RTG-likelihood loss [2305.11340]. The same framework defines an adaptive inference rule through a safe-RTG threshold
\[
\theta_\delta(s)=\max\{r\mid P_{\bar\beta_\theta}(R\ge r\mid s)\ge \delta\},
\]
and a “safe” policy that restricts action selection to returns above that threshold, thereby avoiding OOD queries that yield unpredictable behaviors in vanilla RCRL methods [2305.11340].

“Critic-Guided Decision Transformer” addresses a related inconsistency: the gap between sampled return-to-go inside a single trajectory and expected return across trajectories in stochastic environments [2312.13716]. Its central modification is a learned critic $Q_\phi(R\mid \tau_{0:t-1},s_t,a_t)$ parameterized as a heteroskedastic Gaussian $\mathcal N(\mu_t,\sigma_t^2)$, trained with an asymmetric negative log-likelihood weighted by $\tau_c\in(0,1)$ [2312.13716]. The policy retains the supervised action loss but adds a critic-guidance term based on an expectile loss
\[
\mathcal L_2^{\tau_p}(u)=|\tau_p-\mathbf 1[u<0]|u^2,\qquad
u=\frac{R_t-\mu_t}{\sigma_t},
\]
so that actions whose critic mean falls below the target return are penalized [2312.13716]. The stated objective is to ensure a direct alignment between specified target returns and expected returns of actions, thereby bridging the deterministic nature of RCSL and the probabilistic characteristics of value-based methods [2312.13716].

## 5. Stitching, relabeling, and dynamics shift

A major strand of later work attempts to recover the “stitching” behavior that vanilla RCSL lacks. “How to Provably Improve Return Conditioned Supervised Learning?” introduces the in-distribution optimal return-to-go
\[
f^*(s,h)=\max_{\tau\in T_\beta;\,s_h(\tau)=s} g(\tau,h),
\]
defined as the maximum future return available in the dataset among trajectories visiting state $s$ at step $h$ [2506.08463]. Reinforced RCSL leaves the supervised training loss unchanged, but at test time replaces the raw trajectory return with $\hat f^*(s,h)$, estimated either by scanning the dataset in a tabular or deterministic setting or by regression in a stochastic or large-scale setting [2506.08463]. Its informal Theorem 1 states that if $\pi_f(a\mid s,h)=P_\beta(a\mid s,h,f(s,h))$ for any consistent $f\in F_\beta^C$, then the stitched policy $\pi_\beta^*(a\mid s,h)=P_\beta(a\mid s,h,f^*(s,h))$ satisfies
\[
J(\pi_\beta^*)\ge J(\pi_f),\qquad \forall f\in F_\beta^C,
\]
and a multi-step relabeling scheme yields optimal $Q_h^{*,\beta}(s_h,a_h)$ after $H-1$ passes, so that “RCSL + $k$-step relabeling provably closes the gap to DP methods” [2506.08463].

Model-Based RCSL (MBRCSL) pursues stitching by synthetic rollout rather than return relabeling. It learns a dynamics model $T_\theta(s',r\mid s,a)$ and a behavior policy $\mu_\psi(a\mid s)$ by maximum likelihood from the offline dataset, forward-samples trajectories from $(\mu_\psi,T_\theta)$, retains only those with return larger than the best observed return $g_{\max}$, and then trains a return-conditioned policy on the generated rollout set $\mathcal R$ [2310.19308]. The stated rationale is that the learned model can produce novel transitions not in the dataset, allowing it to stitch pieces of different sub-optimal demonstrations into higher-return trajectories without Bellman backups [2310.19308].

In offline off-dynamics RL, Return Augmented Decision Transformer (RADT) addresses a different stitching problem: source-domain trajectories are useful, but their dynamics differ from the target MDP [2410.23450]. The method augments source-domain returns by aligning the source conditional return distribution with the target one through a transformation $\psi$, written as a KL objective on the source and target return CDFs and, when invertible, as
\[
\psi(g)=G_\beta^{T,-1}(G_\beta^S(g;s,a);s,a).
\]
Its simplified principal theorem states
\[
J^T(\pi^*)-J^T(\hat \pi_f)=O\bigl((N^T+N^S)^{-1/4}\bigr),
\]
matching the same rate obtained without a dynamics shift [2410.23450]. Two practical instantiations are given: RADT-DARA, which relabels each source trajectory with a return built from the original reward plus a transition-density-ratio augmentation, and RADT-MV, which approximates source and target return distributions by Gaussians and performs mean-variance return matching [2410.23450].

## 6. Empirical regimes, applications, and comparative position

The empirical picture across the cited literature is sharply regime-dependent. In simple point-mass environments, return-conditioned methods fail on “Ring-of-fire” because they never observe trajectories with positive $g$ so $\alpha_f\approx 0$; on “Sparse,” they barely reach the goal and are unstable; on “Dense,” they do well because they benefit from a gradient in $g$, although IQL is still slightly better [2206.01079]. On D4RL, antmaze tasks expose both stochastic behavior and the need for trajectory stitching, so RCSL fails while IQL and TD3+BC succeed; HalfCheetah has moderate coverage, so RCSL is decent but dynamic programming remains best; and on pen-human, where state coverage is poor but returns are high, Decision Transformer matches or outperforms dynamic-programming baselines while behavior cloning nearly fails [2206.01079]. The paper’s practical synthesis is explicit: RCSL only works when high-return trajectories are well represented, the environment is near-deterministic or rewards are dense, and stitching is not critical; otherwise DP methods are more robust [2206.01079].

Extensions change that empirical boundary. BR-RCRL reports a Gym-MuJoCo average normalized score of $83.2$ versus Decision Transformer’s $74.7$, and on Atari 10% offline buffers reports Breakout $1239$ versus CQL’s $889$, Q*bert $117$ versus $103$, Pong $138$ versus $131$, and Seaquest with similar scores around $7$ [2305.11340]. CGDT reports that on suboptimal D4RL datasets it outperforms RCSL baselines by $10+$ points and matches or exceeds CQL and IQL, while on a Bernoulli bandit it matches the Bayes-optimal choice for all $p\in\{0.1,\dots,0.5\}$ whereas BC and vanilla Decision Transformer converge to the suboptimal arm [2312.13716]. Reinforced RCSL reports that R-Quantile with $\alpha=0.99$ improves total D4RL Gym score from approximately $490$ or $523$ up to approximately $602$, and raises total AntMaze score from approximately $114$ to approximately $124$ [2506.08463]. MBRCSL reports $91.5\pm 7.1$ return on Point Maze, compared with Decision Transformer’s $57.2\pm 4.1$ and CQL’s $34.8\pm 24.9$, and robotic-task success rates of $0.48\pm 0.04$, $0.51\pm 0.12$, and $0.68\pm 0.09$ on PickPlace, ClosedDrawer, and BlockedDrawer, respectively [2310.19308]. RADT reports that RADT-MV and RADT-DARA often exceed all other methods on mixed $1$T$10$S data and that RADT-MV is consistently the strongest under severe dynamics shifts such as Kinematic and Morphology [2410.23450].

RCRL has also been applied outside standard offline-control benchmarks. Control Transformer reports seen-maze success rates of $84.3\%$ for Point and $90.0\%$ for Ant, and in 20 unseen simulated cluttered Turtlebot3 maps reports $92.5\pm 0.9\%$ success for CT versus $87.6\pm 1.3\%$ for BC-CT and $95.9\pm 0.9\%$ for F-CT; on $7$ real-world start/goal tasks it achieves $3/7$ full successes and $5/7$ partial successes with zero-shot sim-to-real transfer [2211.06407]. In large language model reasoning, an offline RCRL stage built from Expert-Iteration traces conditions on binary “[GOOD]” or “[BAD]” labels attached to partial chains of thought; on GSM8K with SFT initialization, the reported $13$B results are $0.51/0.56/0.69/0.83$ for RCRL, compared with $0.48/0.53/0.68/0.84$ for SFT only, $0.51/0.55/0.67/0.85$ for PPO, and $0.53/0.59/0.71/0.88$ for Expert Iteration [2403.04642]. That study concludes that RCRL is appealingly simple because it reuses off-the-shelf language-model training machinery and a single cross-entropy loss, but that its performance is underwhelming relative to Expert Iteration and PPO on math reasoning tasks, with limited exploration identified as the dominant bottleneck [2403.04642].

Across these lines of work, RCRL occupies a distinctive methodological position. It avoids temporal-difference bootstrapping and the “deadly triad” instabilities emphasized in Reinforced RCSL, and it can avoid Bellman completeness requirements emphasized in MBRCSL [2506.08463][2310.19308]. However, the foundational analysis shows that its coverage and determinism assumptions are strictly stronger than those of classical dynamic-programming-based offline RL, and much of the subsequent literature can be read as attempts to relax exactly those bottlenecks through Bayesian reparameterization, learned critics, in-distribution relabeling, model-based rollout generation, or return-distribution alignment under dynamics shift [2206.01079][2305.11340][2312.13716][2506.08463][2410.23450].

Source: https://www.emergentmind.com/topics/return-conditioned-reinforcement-learning-rcrl