Return-Conditioned Reinforcement Learning

• Return-Conditioned Reinforcement Learning is a framework that uses offline trajectories to condition policies on desired returns through supervised learning.
• It reweights behavior actions based on their likelihood of achieving specified returns, effectively converting policy learning into a conditional prediction task.
• Empirical results in domains like D4RL and robotics demonstrate its promise, though challenges such as trajectory stitching and stochasticity remain.

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 (Brandfonbrener et al., 2022, Kumar et al., 2019). 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 (Brandfonbrener et al., 2022).

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(a\mid s,g)$0

In the infinite-data optimum, if $\pi(a\mid s,g)$1 can represent the true behavior-conditional distribution $\pi(a\mid s,g)$2, the induced policy is

$\pi(a\mid s,g)$3

This form makes explicit that RCRL re-weights behavior actions according to their likelihood of achieving the queried return (Brandfonbrener et al., 2022).

Closely related formulations differ mainly in conditioning granularity. Sequence-modeling variants define

$\pi(a\mid s,g)$4

and train

$\pi(a\mid s,g)$5

by maximum likelihood over offline trajectories (Wang et al., 2023). Other formulations write a time-indexed policy

$\pi(a\mid s,g)$6

with loss

$\pi(a\mid s,g)$7

so that test-time execution uses $\pi(a\mid s,g)$8 (Liu et al., 10 Jun 2025). In the reward-conditioned formulation of behavior policy $\pi(a\mid s,g)$9, the reward-to-go random variable is

$\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$0

and one trains $\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$1 on tuples $\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$2, executing a high target RTG at test time and decrementing it as rewards are observed (Ding et al., 2023).

2. Canonical formulations and architectures

An early formulation appears in “Reward-Conditioned Policies,” which models a single policy

$\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$3

conditioned on a scalar “value target” $\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$4, either the discounted return-to-go or the estimated advantage $\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$5 (Kumar et al., 2019). The paper derives the method from a KL-constrained policy-search objective, obtaining a non-parametric optimum

$\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$6

with the interpretation that, conditioned on achieving return $\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$7, the optimal policy reproduces the data-collection trajectories that achieved $\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$8, while the target distribution over $\pi_\theta(a_t\mid \tau_{0:t-1},s_t,R_t)$9 is shifted toward higher values by $\mathcal D$0 (Kumar et al., 2019). The resulting algorithmic family includes RCP-R and RCP-A, replay-buffer relabeling, optional sample weighting by $\mathcal D$1, and multiplicative conditioning rather than simple concatenation to prevent the network from ignoring $\mathcal D$2 (Kumar et al., 2019).

Transformer instantiations place RCRL inside autoregressive sequence modeling. Decision Transformer-style models tokenize a trajectory into triples $\mathcal D$3, 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 (Wang et al., 2024). In another formulation, the encoder ingests $\mathcal D$4 with positional encodings and causal masking, and predicts $\mathcal D$5 (Wang et al., 2023). More generally, reward conditioning “enables use of generic architectures (MLPs, Transformers, diffusion models) but carries no explicit inductive bias about how different $\mathcal D$6 values relate” (Ding et al., 2023).

Control Transformer extends the sequence-modeling view to long-horizon robot navigation. It represents each navigation trajectory as

$\mathcal D$7

where $\mathcal D$8 is the goal offset and $\mathcal D$9 is the return-to-go, then learns an autoregressive policy

$\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),$0

and a goal-conditioned value function $\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),$1 used at test time to approximate the unknown optimal initial return (Lawson et al., 2022). 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 (Lawson et al., 2022).

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

$\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),$2

and there exists an $\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),$3 such that

$\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),$4

If $\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),$5 and $\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),$6 with $\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),$7, then $\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),$8 (Brandfonbrener et al., 2022). These results show exact recovery in fully deterministic settings, but only under assumptions stronger than those needed by dynamic-programming-based offline RL (Brandfonbrener et al., 2022).

The same analysis gives counterexamples establishing necessity. In the “One-arm” example, $\tau=(o_1,a_1,r_1,\dots,o_H,a_H,r_H),$9, but the $g(\tau)=\sum_{t=1}^H r_t.$0 policy cannot produce return $g(\tau)=\sum_{t=1}^H r_t.$1, so suboptimality $g(\tau)=\sum_{t=1}^H r_t.$2 matches $g(\tau)=\sum_{t=1}^H r_t.$3. In “Max-arm,” conditioning on $g(\tau)=\sum_{t=1}^H r_t.$4 leads RCSL to always pick the red arm, incurring regret approximately $g(\tau)=\sum_{t=1}^H r_t.$5, which shows bias from stochasticity. In “Irremediable bias,” regardless of $g(\tau)=\sum_{t=1}^H r_t.$6, RCSL picks arms $g(\tau)=\sum_{t=1}^H r_t.$7–$g(\tau)=\sum_{t=1}^H r_t.$8, producing irreducible suboptimality (Brandfonbrener et al., 2022). A further negative result is sample-complexity: in deterministic MDPs the required number of samples can be exponential in $g(\tau)=\sum_{t=1}^H r_t.$9 if $s_t$0 (Brandfonbrener et al., 2022).

A distinct theoretical line emphasizes why RCSL is attractive relative to off-policy dynamic programming when function approximation is imperfect. RCSL “never applies $s_t$1; it performs a one-shot supervised fit to the observed $s_t$2 triples,” and therefore avoids the Bellman completeness assumption that DP methods require for stable convergence (Zhou et al., 2023). The “Width Separation” theorem states that there exists a family of deterministic MDPs such that any two-layer ReLU network approximating both $s_t$3 and closed under the Bellman operator requires hidden-layer width $s_t$4, whereas a two-layer ReLU network implementing an RCSL policy $s_t$5 needs width $s_t$6 independent of $s_t$7 (Zhou et al., 2023). 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 (Zhou et al., 2023).

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 (Ding et al., 2023). The paper attributes a core obstacle to “RTG Independence,” meaning that vanilla RCRL feeds $s_t$8 as an input token and treats $s_t$9 and $\pi(a\mid s,g)$0 as independent prediction tasks (Ding et al., 2023). Its proposed decomposition

$\pi(a\mid s,g)$1

leads to an energy-based model with

$\pi(a\mid s,g)$2

trained by a base conditional-likelihood loss plus an auxiliary RTG-likelihood loss (Ding et al., 2023). The same framework defines an adaptive inference rule through a safe-RTG threshold

$\pi(a\mid s,g)$3

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 (Ding et al., 2023).

“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 (Wang et al., 2023). Its central modification is a learned critic $\pi(a\mid s,g)$4 parameterized as a heteroskedastic Gaussian $\pi(a\mid s,g)$5, trained with an asymmetric negative log-likelihood weighted by $\pi(a\mid s,g)$6 (Wang et al., 2023). The policy retains the supervised action loss but adds a critic-guidance term based on an expectile loss

$\pi(a\mid s,g)$7

so that actions whose critic mean falls below the target return are penalized (Wang et al., 2023). 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 (Wang et al., 2023).

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

$\pi(a\mid s,g)$8

defined as the maximum future return available in the dataset among trajectories visiting state $\pi(a\mid s,g)$9 at step $$\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),$$0 (Liu et al., 10 Jun 2025). Reinforced RCSL leaves the supervised training loss unchanged, but at test time replaces the raw trajectory return with $$\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),$$1, estimated either by scanning the dataset in a tabular or deterministic setting or by regression in a stochastic or large-scale setting (Liu et al., 10 Jun 2025). Its informal Theorem 1 states that if $$\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),$$2 for any consistent $$\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),$$3, then the stitched policy $$\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),$$4 satisfies

$$\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),$$5

and a multi-step relabeling scheme yields optimal $$\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),$$6 after $$\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),$$7 passes, so that “RCSL + $$\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),$$8-step relabeling provably closes the gap to DP methods” (Liu et al., 10 Jun 2025).

Model-Based RCSL (MBRCSL) pursues stitching by synthetic rollout rather than return relabeling. It learns a dynamics model $$\hat L(\pi)=-\sum_{\tau\in\mathcal D}\sum_{t=1}^H \log \pi(a_t\mid s_t,g(\tau)),$$9 and a behavior policy $f(s)$0 by maximum likelihood from the offline dataset, forward-samples trajectories from $f(s)$1, retains only those with return larger than the best observed return $f(s)$2, and then trains a return-conditioned policy on the generated rollout set $f(s)$3 (Zhou et al., 2023). 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 (Zhou et al., 2023).

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 (Wang et al., 2024). The method augments source-domain returns by aligning the source conditional return distribution with the target one through a transformation $f(s)$4, written as a KL objective on the source and target return CDFs and, when invertible, as

$f(s)$5

Its simplified principal theorem states

$f(s)$6

matching the same rate obtained without a dynamics shift (Wang et al., 2024). 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 (Wang et al., 2024).

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 $f(s)$7 so $f(s)$8; on “Sparse,” they barely reach the goal and are unstable; on “Dense,” they do well because they benefit from a gradient in $f(s)$9, although IQL is still slightly better (Brandfonbrener et al., 2022). 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 (Brandfonbrener et al., 2022). 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 (Brandfonbrener et al., 2022).

Extensions change that empirical boundary. BR-RCRL reports a Gym-MuJoCo average normalized score of $\pi(a\mid s,g)$00 versus Decision Transformer’s $\pi(a\mid s,g)$01, and on Atari 10% offline buffers reports Breakout $\pi(a\mid s,g)$02 versus CQL’s $\pi(a\mid s,g)$03, Q*bert $\pi(a\mid s,g)$04 versus $\pi(a\mid s,g)$05, Pong $\pi(a\mid s,g)$06 versus $\pi(a\mid s,g)$07, and Seaquest with similar scores around $\pi(a\mid s,g)$08 (Ding et al., 2023). CGDT reports that on suboptimal D4RL datasets it outperforms RCSL baselines by $\pi(a\mid s,g)$09 points and matches or exceeds CQL and IQL, while on a Bernoulli bandit it matches the Bayes-optimal choice for all $\pi(a\mid s,g)$10 whereas BC and vanilla Decision Transformer converge to the suboptimal arm (Wang et al., 2023). Reinforced RCSL reports that R-Quantile with $\pi(a\mid s,g)$11 improves total D4RL Gym score from approximately $\pi(a\mid s,g)$12 or $\pi(a\mid s,g)$13 up to approximately $\pi(a\mid s,g)$14, and raises total AntMaze score from approximately $\pi(a\mid s,g)$15 to approximately $\pi(a\mid s,g)$16 (Liu et al., 10 Jun 2025). MBRCSL reports $\pi(a\mid s,g)$17 return on Point Maze, compared with Decision Transformer’s $\pi(a\mid s,g)$18 and CQL’s $\pi(a\mid s,g)$19, and robotic-task success rates of $\pi(a\mid s,g)$20, $\pi(a\mid s,g)$21, and $\pi(a\mid s,g)$22 on PickPlace, ClosedDrawer, and BlockedDrawer, respectively (Zhou et al., 2023). RADT reports that RADT-MV and RADT-DARA often exceed all other methods on mixed $\pi(a\mid s,g)$23T$\pi(a\mid s,g)$24S data and that RADT-MV is consistently the strongest under severe dynamics shifts such as Kinematic and Morphology (Wang et al., 2024).

RCRL has also been applied outside standard offline-control benchmarks. Control Transformer reports seen-maze success rates of $\pi(a\mid s,g)$25 for Point and $\pi(a\mid s,g)$26 for Ant, and in 20 unseen simulated cluttered Turtlebot3 maps reports $\pi(a\mid s,g)$27 success for CT versus $\pi(a\mid s,g)$28 for BC-CT and $\pi(a\mid s,g)$29 for F-CT; on $\pi(a\mid s,g)$30 real-world start/goal tasks it achieves $\pi(a\mid s,g)$31 full successes and $\pi(a\mid s,g)$32 partial successes with zero-shot sim-to-real transfer (Lawson et al., 2022). In LLM 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 $\pi(a\mid s,g)$33B results are $\pi(a\mid s,g)$34 for RCRL, compared with $\pi(a\mid s,g)$35 for SFT only, $\pi(a\mid s,g)$36 for PPO, and $\pi(a\mid s,g)$37 for Expert Iteration (Havrilla et al., 2024). 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 (Havrilla et al., 2024).

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 (Liu et al., 10 Jun 2025, Zhou et al., 2023). 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 (Brandfonbrener et al., 2022, Ding et al., 2023, Wang et al., 2023, Liu et al., 10 Jun 2025, Wang et al., 2024).