Redistribution-based Cost Inference Improves Sparse Safe Offline RL

AI summary

Overview

Research area: Safe offline reinforcement learning; specifically, temporal credit assignment for cost signals in Constrained Markov Decision Processes (CMDPs).

Technical level: Intermediate. The paper assumes familiarity with MDPs, CMDPs, Lagrangian relaxation, offline RL, and return decomposition, though its central idea is explained in accessible terms.

Scope: The paper introduces Redistribution-based Cost Inference (RCI), a framework that converts trajectory-level binary stop-feedback into dense per-step cost signals for constrained offline policy learning, with theory showing the transformation is lossless and experiments on highway driving and robotic manipulation showing reduced constraint violations.

What This Paper Is About

Safe offline RL methods normally require a cost label on every single timestep of a dataset, but real safety supervisors (human or automated) only produce a single binary signal at the moment they halt an unsafe episode. This leaves a severe credit assignment problem: the dataset says that something went wrong, but not which earlier actions caused it. The paper's goal is to recover dense per-step costs from these sparse stop labels — using return decomposition — and then train a safe offline policy on the recovered costs without sacrificing the theoretical guarantees of the original constrained problem.

Key Contributions

1. The RCI framework, a modular three-stage pipeline: (1) trajectory-level stop-feedback collection, (2) return-decomposition-based cost inference, and (3) constrained offline policy learning. The redistribution stage accepts any return-equivalent decomposition method (instantiated with RUDDER, with GRD noted as directly applicable) and the policy stage accepts any constrained offline RL algorithm (instantiated with BCQ-Lagrangian, with CPQ or CDT noted as substitutable).

2. Theoretical guarantees that the transformation is lossless. The paper defines return-equivalent cost redistribution, proves via a telescoping sum that the redistributed costs exactly reconstruct the episodic cost (Lemma 1), that expected cumulative redistributed cost equals expected episodic cost for any policy (Proposition 1, Constraint Equivalence), and that the feasible policy sets, optimal policy, and Lagrangian saddle point are identical under sparse and redistributed costs (Theorem 1, Policy Invariance, plus Lagrangian Equivalence). Crucially, these hold regardless of the sequence model's prediction accuracy, due to a terminal compensation term.

3. Empirical demonstration on two physical-safety domains — HighwayEnv (highway navigation) and Safe-FetchReach (7-DOF robotic arm reaching with a spherical hazard region) — that RCI substantially lowers violation rates versus sparse and classifier-based baselines without a statistically significant return penalty.

4. Robustness analyses covering heterogeneous dataset compositions (PPO, random, and mixed behavior policies) and two label-corruption regimes (noisy stop-point indices and adversarially flipped labels), plus qualitative spatial analyses of learned cost critics and policy vector fields.

Main Findings

• Lower violation rates on HighwayEnv: On the mixed-composition dataset, RCI cuts violation rate compared to the Sparse and Hazard baselines while maintaining task return. A two-sample t-test against the unconstrained baseline yields t = 0.9962, p = 0.3483, indicating no statistically significant return penalty.

• Roughly fivefold violation reduction across both domains: The conclusion states RCI "reduces violation rates roughly fivefold on two physical-safety domains without statistically significant return penalties." Specific numeric violation-rate values are not reported in the text; the figures present them as bars with standard error.

• Better-shaped spatial cost structure: Learned cost critics visualized over HighwayEnv road geometry show that Sparse concentrates all cost at the terminal unsafe configuration with near-zero values elsewhere; Hazard captures near-terminal danger in a small neighborhood that decays rapidly; RCI produces globally coherent structure, with cost rising gradually as the ego approaches traffic and high-cost regions extending backward along approach corridors.

• Robustness to dataset composition: Violation reduction holds across PPO, Random, and Mixed behavior policies, indicating the redistribution mechanism does not depend on structured or near-optimal exploration (means over 3 seeds, 1,000 episodes per seed).

• Robustness to label noise: Under stop-point shifts drawn uniformly from [−15, 15] in increments of 5, and under adversarial flipping of 20% of trajectory labels, RCI's violation rates respond less sharply to misaligned supervision than Sparse and Hazard, consistent with a smoothing effect of return decomposition. Returns remain stable across noise regimes since termination-time noise alters constraint signals without affecting logged rewards.

• Coherent avoidance in manipulation: On Safe-FetchReach, RCI achieves competitive returns while substantially reducing violations, and the policy vector field shows strong repulsive behavior away from the hazard region under a restrictive budget, demonstrating that spatially coherent avoidance emerges from trajectory-level stop-feedback alone.

• Interpretable budget trade-offs: HighwayEnv trajectories under a restrictive budget (d = P₁₀) maintain large following distances and exit the highway; under a balanced budget (d = P₃₀) they sustain forward progress while respecting distance constraints. Safe-FetchReach vector fields show a smooth degradation from goal pursuit to hazard avoidance as the budget tightens, consistent with an increasing optimal Lagrange multiplier λ*.

Methodology in Plain English

The researchers assume the offline dataset comes with only one bit of safety information per episode: whether it was halted by a safety monitor, and at which step. They start by truncating each trajectory at its first violation and assigning an episodic cost of 1 (or 0 if the trajectory is safe).

Then they train a sequence model (an LSTM, in this instantiation) to predict that episodic cost from the trajectory prefix, using squared error. The per-step cost is obtained by taking differences between consecutive predictions — if the predicted chance of eventual failure jumps sharply at a timestep, that transition matters. A final compensation term at the terminal step nudges the summed per-step costs to exactly match the original trajectory-level label, so the transformation introduces no systematic bias even if the sequence model is inaccurate.

The resulting cost-augmented dataset feeds a standard constrained offline RL algorithm (BCQ-Lagrangian), which uses a VAE-based behavior cloning constraint against distributional shift and an adaptive Lagrange multiplier updated from the batch-average cost relative to the budget. The Bellman target subtracts λ times the inferred cost from the discounted value.

For evaluation, the authors generate 5,000 offline episodes per environment using PPO-trained behavioral policies that ignore safety. They compare against three baselines sharing the same BCQ-Lagrangian architecture: Reward-Only (budget set to infinity), Sparse (raw terminal cost label, no redistribution), and Hazard (a two-head binary classifier summing a "appears in an unsafe trajectory" head and an "unsafe termination point" head, trained with focal loss for class imbalance). They sweep the safety budget over the 10th–50th percentiles of the dataset's trajectory cost distribution in increments of 10, train three independent policies per configuration, and evaluate the lowest-violation policy over 1,000 online episodes against ground-truth safety labels.

Why This Matters

Impact on research. The paper relaxes a nearly universal assumption in safe offline RL — that dense per-step cost annotations are available — and shows, both theoretically and empirically, that a far weaker supervision signal suffices. It also carries the return-decomposition principle from sparse-reward RL into the cost and constraint domain, establishing preservation results for the feasible policy set and the Lagrangian saddle point that make the transfer principled rather than merely heuristic. Prior sparse-supervision work (TraCeS, RLSF) assumes the labeler can be queried during training; this work targets the strictly offline setting where no post-hoc queries are possible.

Real-world applications:

• Autonomous driving, where an automated monitor issues a single stop command at the first detected violation rather than scoring every frame.
• Robotic manipulation around designated hazard zones, where a supervisor halts the arm upon entering a forbidden region.
• Clinical decision support, cited by the paper as a high-stakes domain for safe RL, where per-step cost annotation is impractical.
• Any deployed physical AI system that already produces halt events as a byproduct of its safety monitoring, allowing existing logs to be reused as safety training data.

Industry relevance. The supervision modality RCI targets — "stop-feedback" issued by automated safety systems or human overseers at the point of intervention — is what physical AI deployments already produce. This means organizations do not need to build new dense-annotation pipelines to train constrained policies; they can reuse intervention logs. The modular design also makes adoption easier, since practitioners can slot in their own decomposition method and their own constrained offline RL algorithm.

Future Directions

• Coverage and conservatism: The authors note RCI inherits standard offline RL coverage requirements — datasets skewed toward unsafe trajectories risk overly conservative policies, and sparse safe coverage may leave policy optimization underspecified. How to diagnose and mitigate this is left open.
• Statistical versus causal attribution: The redistribution mechanism captures statistical associations between trajectory prefixes and violations rather than causal hazard. Closing that gap — identifying causally hazardous actions rather than merely risky situations — is an explicit open direction.
• Multi-constraint extension: RCI is currently restricted to a single binary constraint. Per-channel decomposition is described as straightforward, but balancing competing objectives raises open questions.
• Richer supervision and architectures: The authors list graduated severity feedback (instead of a single binary stop signal), uncertainty-aware redistribution under distributional shift, and more expressive sequence architectures as future work.

Target Audience

Researchers and practitioners in safe reinforcement learning, offline RL, and constrained optimization who need to train safe policies from datasets that lack dense per-step cost labels. It is also relevant to applied roboticists and autonomous-systems engineers whose deployment logs contain only intervention or halt events, and to readers interested in temporal credit assignment and return decomposition. Readers without a background in CMDPs and Lagrangian methods should expect to engage with the appendix proofs and the Bellman-target formulation.

Authors’ abstract

Safe offline RL typically assumes access to dense per-step cost annotations, but in practice supervisors provide only trajectory-level stop-feedback: a binary signal at the first unsafe transition, with no per-step attribution. We frame this as a temporal credit assignment problem and propose the Redistribution-based Cost Inference (RCI) framework, which converts sparse stop-feedback into dense per-step costs via return decomposition, then trains a constrained offline policy on the augmented dataset. We show that return-equivalent redistribution preserves the feasible policy set and the optimal Lagrangian in a CMDP, establishing that the transformation is lossless in theory while yielding better-conditioned cost critic learning in practice. Experiments on highway driving and robotic manipulation demonstrate substantially lower violation rates than sparse and classifier-based baselines, with robustness to heterogeneous dataset compositions and label noise.