PaperlayerView exact revision on arXiv

Source-linked AI summary

Q-learning Penalized Transformer for Safe Offline Reinforcement Learning

Shengchao Hu, Peng Wang, Jifeng Hu, Qiyang Zhou, Anning Hu, Li Shen, Ya Zhang, Dacheng Tao

arXiv:2609.34426v2cs.LG

TL;DR

Safe offline RL must balance reward, constraint satisfaction, and behavior regularization, while offline cost-estimation bias can make constrained optimization challenging. QPT couples a conditional Transformer with reward and cost Q-functions for training and inference. Across 38 DSRL tasks it outperforms strong safe offline RL baselines and adapts zero-shot to different constraint thresholds.

  • Problem

    Safe offline RL must balance reward maximization and constraint satisfaction while offline policy learning faces cost-estimation bias that can make policies unsafe or overly conservative.

    Source: §1 · p. 1, §3.1 · p. 3

  • Method

    QPT trains a conditional Transformer with a Q-shaped reward-and-cost objective, then reuses the learned Q-functions to select the highest-reward feasible action at inference.

    Source: §1 · p. 2

  • Results

    Across 38 DSRL tasks, QPT consistently outperforms strong safe offline RL baselines in reward and constraint satisfaction and shows robust zero-shot adaptation to varying thresholds.

    Source: §5 · p. 8

  • Takeaways & Limitations

    QPT's reported zero-shot adaptation supports using one trained policy across different constraint thresholds.

    Source: §5 · p. 8

    The theoretical analysis is scoped to MDPs with binary rewards and costs and the stated assumptions.

    Source: §3.5 · p. 5

Abstract

from arXiv · show

This paper addresses the problem of safe offline reinforcement learning, which involves training a policy to satisfy safety constraints using an offline dataset. This problem is inherently challenging as it requires balancing three highly interconnected and competing objectives: satisfying safety constraints, maximizing rewards, and adhering to the behavior regularization imposed by the offline dataset. To tackle this trilogy challenge, we propose Q-learning Penalized Transformer policy (QPT), a \emph{training--inference consistent} framework that bridges conditional sequence modeling with constraint-aware value estimation. QPT trains a Transformer policy that generates actions conditioned on trajectory context and target return/cost, retaining strong behavior regularization. To inject explicit safety semantics during learning, we augment sequence-model training with a Q-shaped penalty using learned reward and cost Q-functions to favor high return under low constraint violation. At inference, the same Q-functions enforce the cost threshold and choose the highest-reward feasible action, closing the loop between training and deployment. We provide a principled analysis under stylized near-deterministic CMDPs, characterizing how Q-penalized conditional generation improve safety and performance. Empirically, QPT consistently outperforms strong safe offline RL baselines across 38 tasks on the DSRL benchmark, and exhibits robust zero-shot adaptation to different constraint thresholds.

1. Introduction

Safe offline RL must balance safety constraints, reward maximization, and behavior regularization, which can conflict in offline data. QPT combines conditional sequence modeling with learned reward and cost Q-functions for training and inference.

1. Introduction

  • Safe offline RL must balance constraint satisfaction, reward maximization, and offline policy regularization, whose interdependencies make joint optimization difficult.These objectives can conflict when offline data includes unsafe trajectories or when excluding unsafe trajectories omits high-reward data.Source: §1 · pp. 1–2
  • Unsafe trajectories in the dataset can make behavior regularization compromise constraint satisfaction, while excluding them can discard critical high-reward data.Source: §1 · pp. 1–2
  • QPT conditions Transformer-generated actions on trajectory context and target return/cost specifications, retaining behavior regularization while adding a Q-shaped reward-and-cost penalty.Learned reward and cost Q-functions bias generation toward high return with low constraint violation.Source: §1 · p. 2
  • At inference, QPT reuses the same Q-functions to test feasibility under a cost threshold and select the highest-reward feasible action.Source: §1 · p. 2
  • QPT is analyzed in stylized near-deterministic CMDPs and evaluated on 38 DSRL tasks, where it consistently outperforms strong safe offline RL baselines and adapts zero-shot to varying thresholds.Source: §1 · p. 2

2. Related Work

Prior work addresses offline distribution shift through behavior regularization and sequential modeling, while safe RL uses constrained optimization, action correction, and model-based methods. Safe offline RL combines these lines of work, and QPT frames reward and cost control as controllable loss terms.

2. Related Work

  • Offline policy deviation from the behavior policy can cause performance degradation, motivating regularized or constrained dynamic programming.Source: §2 · p. 2
  • Conditional sequence models predict future actions from past experience, constrain policies within behavior boundaries, and enable zero-shot adaptability.Source: §2 · p. 2
  • Safe RL approaches include primal-dual optimization, correction-based projection of unsafe actions, and model-based methods that can increase computational complexity.Source: §2 · p. 2
  • Safe offline RL includes DICE-style and Lagrangian methods, sequential modeling such as CDT, diffusion approaches such as TREBI and FISOR, and data-centric methods.QPT contrasts with these approaches by treating reward maximization and cost control as controllable loss terms.Source: §2 · p. 2

3. Methodology

QPT addresses safe offline RL by combining a trajectory-conditioned Transformer with reward and cost Q-functions, using Q-based penalties in training and the same value estimates to guide deployment. Its theoretical analysis reports improved expected reward and reduced expected cost over standard sequence-model training under a stylized setting.

3.1. Problem Setup

  • QPT frames learning around maximizing expected reward while satisfying a cost threshold, addressing offline value-estimation bias and the difficulty of adapting policies to new thresholds.Sequence-modeling methods may imitate unsafe data, while filtering unsafe trajectories can remove high-reward transitions.Source: §3.1 · p. 3

3.2. Conditional Transformer Policy

  • The conditional Transformer predicts actions from recent trajectory context, including state, action, return-to-go, and cost-to-go tokens.Training is autoregressive over the most recent K steps, and the state-token prediction head uses mean squared error for continuous actions.Source: §3.2 · p. 3

3.3. Q-learning Penalization

  • Q-learning penalization combines action-prediction MSE with reward-maximizing and cost-sensitive Q terms, with weights balancing the losses.Reward and cost Q-networks are learned using Bellman evaluation and target-policy actions.Source: §3.3 · p. 4

3.4. Data Augmentation and Ensemble

  • For infeasible target pairs, QPT augments training with the safest trajectory achieving maximum reward and uses Q-networks to select high-return, low-cost candidates in parallel.The candidate search can incur extra computation and greater susceptibility to noisy or suboptimal Q estimates as the candidate set grows.Source: §3.4 · pp. 4–5

3.5. Theoretical Analysis

  • Under mild coverage and regularity assumptions in a stylized setting, Q-penalized training improves expected reward and reduces expected cost relative to standard sequence-model training.The theorem’s setting includes near-deterministic transitions, rewards, and costs.Source: §3.5 · p. 5

4. Experiment

Across the DSRL experiments, QPT is reported as the only method that consistently meets safety requirements while achieving the highest returns in most cases. Ablations show how Q-learning penalties, data augmentation, ensembling, and target settings affect reward and cost.

4. Experiment

  • Evaluations span Safety-Gymnasium, Bullet-Safety-Gym, and MetaDrive tasks, using normalized reward and cost with safety prioritized and a cost limit of 10.The reported Table 1 values average 20 evaluation episodes and 3 random seeds.Source: §4 · p. 5, Table 1 · p. 6
  • QPT is the only method reported to consistently achieve satisfactory safety across tasks while attaining the highest returns in most cases.The comparison uses normalized reward and cost, with a cost threshold of 1.Source: §4 · pp. 6–7, Table 1 · p. 6

4.1. Ablation

  • Adding reward and cost Q-network penalization improves reward and cost, ensembling further boosts performance, and data augmentation improves cost by stitching safe trajectories into training data.The reported comparisons are Exp 2 vs. 3/4, Exp 2 vs. 6 and Exp 5 vs. 8, and Exp 1 vs. 2 and Exp 7 vs. 8, respectively.Source: §4.1 · p. 7, Table 2 · p. 7
  • In HopperVelocity, increasing η1 raises reward only within a range and can sharply reduce it beyond η1 = 1; increasing η2 lowers reward while reducing cost within a range.The hyperparameter study covers SwimmerVelocity and HopperVelocity.Source: §4.1 · p. 7, Figure 1 · p. 8
  • When both methods satisfy the constraint, QPT consistently outperforms CDT; in CarButton1 and PointButton1, CDT’s higher rewards come with greater safety violations.Both methods improve when conditioned on a higher cost threshold; Figure 2 plots reward and cost across target-cost settings.Source: §4.1 · p. 8, Figure 2 · p. 8

5. Conclusion

QPT couples a conditional Transformer policy with reward and cost Q-functions to balance behavior regularization, reward, and constraint violation.

5. Conclusion

  • Across 38 DSRL tasks, QPT consistently outperforms strong safe offline RL baselines and robustly adapts zero-shot to different constraint thresholds.The paper also provides theoretical analysis in a stylized setting.Source: §5 · p. 8
  • QPT combines a conditional Transformer policy with reward and cost Q-functions through a Q-shaped objective.The objective favors high return under low constraint violation while retaining behavior-regularized learning.Source: §5 · p. 8
  • The paper analyzes QPT theoretically in a stylized setting.Source: §5 · p. 8

Impact Statement

The paper concerns safe policies learned from offline datasets and identifies no societal consequences requiring specific emphasis.

Impact Statement

  • The work concerns safe policies learned from offline datasets, and the authors identify no societal consequences requiring specific emphasis.Source: Impact Statement · p. 9

A. Proof of Theorem 3.1

The proof analyzes Q-penalized updates under stylized assumptions and contrasts their reward- and cost-side guarantees with a baseline objective.

A. Proof of Theorem 3.1

  • The analysis is explanatory, not a full characterization of deep offline RL with function approximation, and assumes binary signals, near-determinism, and consistent conditioning.The theorem adds improvement terms for the probabilities of selecting reward-maximizing and cost-minimizing actions.Source: A. Proof of Theorem 3.1 · p. 13
  • Under stated reward-side conditions, the corollary equates the returns of the optimal, baseline-induced, and QPT-induced policies.The cost-side result analogously equates costs under its corresponding conditions.Source: A. Proof of Theorem 3.1 · p. 13
  • Reward and cost conditions are analyzed separately because maximizing reward and minimizing cost may not be simultaneously achievable under one conditioning specification.QPT combines reward/cost terms during training and enforces a user-specified cost threshold at inference.Source: A. Proof of Theorem 3.1 · p. 13

B.1. Algorithm Pseudocode

  • Algorithm 1 summarizes QPT’s pipeline, including Q-network and policy updates during training and candidate-action selection at inference.Source: A. Proof of Theorem 3.1 · p. 14

B.2. Data Augmentation

Data augmentation relabels trajectories to support imitation of a rewarding, safe option when the requested return-cost pair is infeasible.

B.2. Data Augmentation

  • When a desired return-cost pair is infeasible, augmentation relabels its associated Pareto trajectory’s returns so the agent can imitate the most rewarding and safe trajectory.In the passage, the Reward Frontier is the maximum reward at cost κ among costs represented in dataset D.Source: §B.2 · p. 14
  • Figure 3 illustrates the augmentation by mapping Pareto-optimal trajectories to their corresponding augmented return-cost pairs.Source: §B.2 · p. 14

B.3. Ensemble

The ensemble generates actions from multiple return–cost conditioning pairs, then uses learned Q-networks to select an action under reward and safety criteria.

B.3. Ensemble

  • Training RTG and CTG values are cumulative discounted rewards and costs from each state to the terminal state; inference starts from the DSRL benchmark’s default pair.The training values are derived from observed trajectories, while the default inference pair serves as the baseline for candidate generation.Source: §B.3 · p. 14
  • Candidate RTG values are perturbed while CTG stays fixed, searching for high reward under a predetermined cost threshold.The asymmetric perturbation follows the objective of maximizing expected returns while adhering to a fixed cost constraint.Source: §B.3 · p. 15
  • All candidate pairs are batched through the Transformer, and learned reward and cost Q-networks evaluate the resulting actions for selection.Selection can prioritize reward subject to hard safety constraints or use a parameterized reward–safety trade-off.Source: §B.3 · p. 15

C.1. Environment descriptions

The appendix describes benchmark environments and tasks, a trajectory relabeling procedure, and QPT’s experimental hyperparameter setup.

C.1. Environment descriptions

  • Safety-Gymnasium covers hazard-avoidance tasks for Car and other agents, while Bullet-Safety-Gym evaluates four agents on Circle and Run tasks.Safety-Gymnasium includes Button, Push, and Goal tasks at two difficulty levels; Bullet-Safety-Gym uses shorter, more straightforward tasks.Source: §C.1 · pp. 15–16
  • The augmentation procedure samples costs within the dataset range, samples rewards above the reward frontier, selects a reward-maximizing trajectory within the sampled cost, then relabels and appends it.The selected trajectory is constrained by C(τ) ≤ κ_i.Source: §C.1 · p. 16
  • MetaDrive varies road complexity and traffic density to test offline RL in self-driving scenarios that mimic real-world driving challenges.Its road conditions are easy, medium, or hard, and traffic density is sparse, medium, or dense.Source: §C.1 · p. 16
  • The benchmark overview covers varied safe-RL tasks and complexities, while Figure 4 visualizes representative environments and tasks.Table 3 is identified as an overview of benchmarks and tasks used for dataset collection.Source: §C.1 · pp. 15–16, Figure 4 · p. 16, Table 3 · p. 17

C.2. Hyperparameters

  • Reward and cost Q-networks use four linear Mish layers, and experiments use 105 gradient steps except on MetaDrive with rollouts at maximum episode length.Table 4 lists QPT hyperparameters; the passage does not specify the exception’s gradient-step count.Source: §C.2 · pp. 16–17, Table 4 · p. 17

C.3. Ablation of the number of candidate actions in Ensemble

The candidate-action ablation finds that increasing ensemble size helps initially but can eventually plateau or reduce performance, with added computational and estimation risks.

C.3. Ablation of the number of candidate actions in Ensemble

  • Increasing candidate count initially improves performance, but the relationship is non-monotonic and performance eventually plateaus or declines.The ablation uses MetaDrive harddense; Table 5 reports the impact of candidate target reward and cost pairs.Source: §C.3 · p. 17, Table 5 · p. 17
  • QPT uses 50 candidate actions by default during inference.The ablation varies the number of sampled candidates in MetaDrive harddense.Source: §C.3 · p. 17
  • Larger candidate sets increase computational cost and susceptibility to noisy or suboptimal pairs from biased Q-network estimates.A broader search may improve performance, but the cited passage does not claim that it always does.Source: §C.3 · p. 17
Loading 2609.34426v2…