Source-linked AI summary

Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling

Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, Jun Zhu

arXiv:2209.14548v2cs.LG

TL;DR

Offline RL weighted regression can still select unseen actions because common policy models lack sufficient distributional expressivity. The paper decouples behavior generation from action evaluation, using expressive diffusion models and in-sample planning. On D4RL, the approach outperforms state-of-the-art algorithms in most tasks and learns from heterogeneous behavior datasets.

  • Problem

    Weighted regression may select unseen actions because conventional policy models have limited distributional expressivity, especially for diverse behaviors.

  • Method

    The method decouples policy learning into an expressive generative behavior model and an action evaluation model, using diffusion modeling and in-sample planning.

  • Results

    The approach outperforms state-of-the-art algorithms in most D4RL tasks and supports learning from heterogeneous datasets while avoiding out-of-sample actions.

  • Takeaways & Limitations

    High-capacity generative behavior models can support diverse behavior learning in offline RL while preserving in-sample action selection.

  • Takeaways & Limitations

    The approach depends on high-fidelity behavior modeling, while weighted regression is constrained by the need for exact, derivable policy densities.

Abstract

from arXiv · show

In offline reinforcement learning, weighted regression is a common method to ensure the learned policy stays close to the behavior policy and to prevent selecting out-of-sample actions. In this work, we show that due to the limited distributional expressivity of policy models, previous methods might still select unseen actions during training, which deviates from their initial motivation. To address this problem, we adopt a generative approach by decoupling the learned policy into two parts: an expressive generative behavior model and an action evaluation model. The key insight is that such decoupling avoids learning an explicitly parameterized policy model with a closed-form expression. Directly learning the behavior policy allows us to leverage existing advances in generative modeling, such as diffusion-based methods, to model diverse behaviors. As for action evaluation, we combine our method with an in-sample planning technique to further avoid selecting out-of-sample actions and increase computational efficiency. Experimental results on D4RL datasets show that our proposed method achieves competitive or superior performance compared with state-of-the-art offline RL methods, especially in complex tasks such as AntMaze. We also empirically demonstrate that our method can successfully learn from a heterogeneous dataset containing multiple distinctive but similarly successful strategies, whereas previous unimodal policies fail.

1 INTRODUCTION

Offline RL methods use weighted regression to stay near dataset behavior, but limited policy expressivity can still select unseen actions. The paper addresses this by separating expressive behavior modeling from action evaluation and reports strong benchmark performance.

  • Offline RL is difficult because incomplete datasets make out-of-sample actions unreliable during policy evaluation.
  • Weighted regression is intended to prevent querying actions absent from the dataset, but this motivation can fail in some settings.
  • Unimodal Gaussian policy models cannot represent diverse behaviors and may assign probability to low-density actions between modes.
  • Exact, derivable density requirements make replacing these restricted distributions with more expressive classes nontrivial.
  • SfBC decouples policy learning into a generative behavior model and action evaluation, selecting among behavior-model candidates using importance weights.
  • The method uses diffusion models for diverse behavior modeling and combines them with planning-based action evaluation; experiments show competitive or superior performance, especially on AntMaze.

2 BACKGROUND

Offline policy optimization replaces unavailable on-policy state visitation with behavior-distribution data and constrains learned policies toward dataset support. Weighted regression requires tractable policy densities, while diffusion models learn distributions through noising and reversal.

  • An MDP comprises state and action spaces, transition and reward functions, and a discount factor, with the goal of maximizing expected discounted return.
  • Without online collection, estimating a policy’s discounted state visitation is difficult, motivating surrogate objectives based on behavior-policy visitation.
  • Offline methods constrain the learned policy toward behavior-policy support while maximizing learned action values.
  • Weighted regression projects an improved policy onto a parameterized policy using exp(αQφ(s, a)) as regression weights.
  • Weighted regression requires an exact policy density, restricting πθ to distribution classes with tractable density expressions and limiting expressivity.
  • Diffusion models add noise to a data distribution through a forward process and learn its reverse transformation.
  • Diffusion sampling starts from Gaussian noise and numerically solves the reverse process to produce a final sample.

3 METHOD

SfBC addresses limited policy expressivity by separating behavior modeling from action evaluation, using expressive generative models to represent diverse behaviors and selecting among behavior-generated candidates.

  • Learning from Diverse Behaviors: Conventional weighted regression can project diverse optimal policies onto narrow unimodal distributions, potentially selecting low-density or unseen actions.A bimodal target fitted by a unimodal model covers the low-density region between its modes.
  • Learning from Diverse Behaviors: The Bidirectional-Car task illustrates that unimodal policies may stay between two similarly rewarding directions, whereas expressive policies can preserve either behavior.The task contains leftward and rightward optimal actions with equal reward.
  • Learning from Diverse Behaviors: Weighted regression restricts practical distribution choices because it requires exact, derivable action densities and may assume unknown behavior or optimal-policy forms.These requirements make adopting more expressive distribution classes nontrivial.
  • Selecting from Behavior Candidates: SfBC decouples policy learning into an expressive generative behavior model and an action evaluation model instead of directly parameterizing the policy.Behavior samples are available from the offline dataset, making behavior modeling more tractable than directly sampling the optimal policy.
  • Selecting from Behavior Candidates: The policy is proportional to behavior density times exp(αQφ), with α controlling the trade-off between conservative behavior matching and greedy improvement.As α approaches zero, the policy returns to behavior; as α grows, it approaches a greedy policy.
  • Selecting from Behavior Candidates: SfBC samples M actions from the learned behavior model, evaluates them with Qφ, and resamples using exp(αQφ) as weights.This procedure directly uses behavior-generated candidates and is described as analogous to rejection sampling.

4 PRACTICAL IMPLEMENTATION

The practical implementation combines diffusion-based behavior modeling with planning-based Q-learning. Implicit in-sample planning avoids bootstrapping over unseen actions while stitching together successful dataset subtrajectories.

  • Behavior Modeling: High-fidelity behavior modeling is necessary because generating out-of-sample actions causes extrapolation error, while missing in-sample actions restricts the feasible policy space.Gaussian models lack expressivity, and VAEs present separate behavior-modeling challenges.
  • Behavior Modeling: SfBC uses a state-conditioned diffusion model trained to denoise perturbed behavior actions and generate samples from the learned behavior distribution.The model predicts time-dependent noise, and inverse-ODE sampling transforms Gaussian noise into behavior actions.
  • Q-Learning via In-Sample Planning: One-step Bellman bootstrapping is computationally inefficient and can accumulate estimation errors over long horizons, with diffusion sampling adding further cost.The dependence on many extrapolation steps is especially costly when policy actions require iterative sampling.
  • Q-Learning via In-Sample Planning: Implicit planning iteratively stitches successful subtrajectories and provides optimistic action evaluations in AntMaze.The visualized Q-targets are updated across consecutive value iterations.
  • Q-Learning via In-Sample Planning: The planning operator combines the fast contraction of an n-step behavior operator with the more desirable fixed point of the policy operator.Its fixed point is proven convergent and bounded between Qπ and Q∗.
  • Q-Learning via In-Sample Planning: Recursive targets are computed from dataset trajectories and alternated with fitting the action evaluation model Qφ.The scheme avoids bootstrapping over unseen actions and generalizes actions across similar states by stitching trajectories.

5 RELATED WORK

Related offline RL methods reduce extrapolation error through policy regularization, pessimistic values, weighted regression, or dataset-only Q-learning. SfBC instead emphasizes expressive behavior modeling and in-sample planning for long-horizon tasks.

  • Reducing extrapolation Error in Offline RL: Offline RL methods commonly trade off return maximization against staying near behavior data to limit extrapolation error.Approaches include policy regularization, value pessimism, weighted regression, and Q-learning without policy-generated actions.
  • Dynamic Programming over Long Horizons: Behavior-Q policy extraction can avoid dynamic-programming extrapolation but performs poorly when tasks require stitching successful subtrajectories.Maze-like environments may require hundreds of steps before rewards are reached, allowing discounting and estimation error to accumulate.
  • Generative Models for Behavior Modeling: Table 1 reports normalized SfBC evaluation scores across D4RL tasks, highlighting results within 5 percent of each task’s maximum.The reported experiments use 15 diffusion steps.
  • Generative Models for Behavior Modeling: Generative behavior models such as Gaussians and VAEs have been used mainly as explicit distributional constraints for another policy model.SfBC differs by directly using the learned behavior model to generate action candidates.

6 EXPERIMENTS

SfBC is evaluated on D4RL benchmarks, diverse-behavior tasks, and ablations of planning and behavior modeling. It performs especially strongly on complex AntMaze tasks and heterogeneous datasets, where expressive behavior modeling and implicit planning are most relevant.

  • Evaluations on D4RL Benchmarks: SfBC outperforms most existing methods by large margins on sparse-reward AntMaze tasks.Both SfBC and VEM use implicit in-sample planning, suggesting episodic planning benefits stitching in maze-like environments.
  • Evaluations on D4RL Benchmarks: SfBC remains highly competitive on easier MuJoCo locomotion tasks, with larger gains on mixed-policy Medium-Expert data than uniform Medium data.This pattern matches the method’s motivation to model diverse behaviors.
  • Learning from Diverse Behaviors: In Bidirectional-Car, SfBC is the only tested method that consistently reaches either endpoint within the rated time.Other methods solve the Single-Side task, but fail when both high-reward directions are present.
  • Learning from Diverse Behaviors: SfBC selects high-absolute-value actions across nearly all states, whereas unimodal actors fail to select either optimal action when two rewarding options exist.The task’s optimal actions are 1.0 and -1.0, and white space in the visualization denotes suboptimal decisions.
  • Ablation Studies: Diffusion-based behavior modeling outperforms Gaussian and VAE variants by a clear margin in most experiments, especially on heterogeneous Medium-Expert data.The variants share the same training framework and differ only in behavior-modeling method.
  • Ablation Studies: Implicit planning helps on AntMaze-Medium and AntMaze-Large but is less important for MuJoCo locomotion tasks.The ablation compares value-iteration numbers K in {1, 2, 3, 4, 5}, with K = 1 using vanilla returns R(0).

7 CONCLUSION

The paper concludes that SfBC addresses limited policy expressivity by separating behavior modeling from action evaluation. Diffusion-based behavior modeling and planning-based Q-learning together support heterogeneous continuous-action learning while reducing out-of-sample action selection.

  • Conclusion: SfBC decouples policy learning into a behavior model and an action evaluation model to address limited policy expressivity.This design lets the policy inherit distributional expressivity from a diffusion-based behavior model.
  • Conclusion: The method combines high-fidelity diffusion behavior modeling with a planning-based operator to reduce extrapolation error.Training alternates between action-evaluation learning and recursively updated planning targets.
  • Conclusion: SfBC enables learning from heterogeneous continuous-action datasets while avoiding selection of out-of-sample actions.Its behavior candidates are sampled from a learned behavior policy and evaluated by the action-evaluation model.
  • Conclusion: On D4RL, SfBC outperforms state-of-the-art algorithms in most tasks.The conclusion presents this as the paper’s overall experimental outcome.

B.1 IMPLEMENTATION DETAILS OF SFBC

SfBC uses a score-based behavior model, a separate Q-value model, and iterative in-sample planning. Implementation choices trade diffusion fidelity and evaluation cost against runtime, with task-specific value-iteration settings.

  • Network Architecture: SfBC uses a conditional score-based behavior model and a separate action-evaluation model that outputs Q-values for state-action pairs.The behavior architecture resembles a dense U-Net, while the evaluator is a 2-layer MLP with 256 hidden units and SiLU activations.
  • Behavior Training: The score-based model is trained with Adam, batch size 4096, 500 data epochs, and learning rate 1e-4.Behavior perturbations use a variance-preserving SDE with βmin = 0.1 and βmax = 20.
  • Action Evaluation via In-Sample Planning: Action evaluation uses iterative planning with K = 2 for MuJoCo, K = 4 for AntMaze-umaze, and K = 5 for other AntMaze tasks.Each iteration recalculates Q-targets using the latest policy, with inverse temperature α = 20 and 16 Monte Carlo samples.
  • Diffusion Sampling: Diffusion sampling uses D = 15 steps for reported Table 1 results, while 10–25 steps generally perform similarly in MuJoCo and 15–25 in AntMaze.A specialized third-order diffusion ODE solver reduces the steps relative to ordinary RK45 sampling.
  • Evaluation: Evaluation selects the highest-Q action among M behavior candidates, and experiments average results over multiple trials and test seeds.AntMaze and MuJoCo use 4 random-seed trials; other experiments use 3.
  • Runtime: Training the behavior model takes 10.5 hours for 600 epochs, while action-evaluation training costs about 31 minutes per value iteration on an RTX 2080Ti.For halfcheetah-medium with K = 5, action evaluation takes about 155 minutes in addition to behavior-model training.
  • Runtime: Parallel evaluation reduces the diffusion model’s theoretical inference overhead by evaluating multiple test seeds simultaneously on GPUs.SfBC otherwise requires at least D times the network-inference time of non-diffusion methods.

C THEORETICAL ANALYSIS

The planning-based operator Tπ_μ is analyzed through monotonicity, contraction, and fixed-point bounds. Its error contraction can be especially rapid when initial Q estimates are pessimistic.

  • Tπ_μ is monotone and at least a γ-contraction.These properties support convergence of the planning-based operator.
  • The fixed point of Tπ_μ lies between Qπ and Q∗ for every state-action pair.The bound is stated as Qπ(s, a) ≤ Q̃(s, a) ≤ Q∗(s, a).
  • When current Q estimates are strongly pessimistic, Tπ_μ provides a relatively optimistic update and contracts estimation error at roughly γ^n∗(s,a).This regime commonly occurs early in training when initial Q estimates are near zero.
  • The operator’s selected iteration n∗(s,a) balances repeated applications of Tμ with the policy evaluation operator Tπ.n∗(s,a) is defined by maximizing the resulting evaluated Q estimate over 0≤n≤N.
  • SfBC is compared with policy regression, return-weighted regression, policy-regularized off-policy, and sequence-modeling baselines.The listed examples include IQL, VEM, AWR, BAIL, BCQ, and CQL.

F TRAINING CURVES

Figure 8 reports SfBC training curves across MuJoCo and AntMaze tasks while varying diffusion steps, with results aggregated over four random seeds.

  • Training curves compare SfBC across different diffusion-step settings on MuJoCo and AntMaze tasks.The figure reports mean and standard deviation over four seeds.

G.1 VEM

The section relates SfBC’s in-sample planning operator to VEM and EMaQ while distinguishing their assumptions and roles. SfBC uses behavior candidates for Q-target estimation and avoids VEM’s deterministic-environment requirement.

  • VEM: VEM’s expectile parameter τ interpolates the Bellman expectation operator at τ=0.5 and the Bellman optimality operator as τ approaches 1.0.Its operator uses different weights depending on whether the target exceeds V(s).
  • VEM: VEM and its related multi-step operator share the same fixed point, supported by the expectile operator’s monotonicity.The monotonicity increases with τ.
  • VEM: VEM assumes deterministic rewards and transitions because its update compares scalar V(s) with a random next-state target in stochastic environments.The deterministic assumption treats r(s,a) and P(·|s,a) as Dirac.
  • SfBC versus VEM: SfBC instead uses an in-sample planning Q-operator with a hypothetically improved policy for optimistic planning.Its comparison with VEM highlights a different planning construction.
  • Theoretical distinction: Replacing Tμ with Tπ_μ removes the monotonicity property, although the fixed point remains bounded between Qπ and Q∗.The bound is retained despite this theoretical trade-off.
  • SfBC versus EMaQ: Unlike EMaQ’s interpolation hyperparameter, SfBC’s N denotes the number of Monte Carlo behavior-policy samples used to estimate training Q-targets.For EMaQ, N=1 corresponds to Tμ and N→∞ approaches T∗.
Loading 2209.14548v2…