Redistribution-based Cost Inference Improves Sparse Safe Offline RL

Problem Setting§

Safe offline RL typically requires dense per-step cost labels. In practice, only sparse stop-feedback is available: a binary signal at the first unsafe transition, with no attribution of how much each step contributed to the danger. This is a temporal credit assignment problem.

RCI Framework§

The Redistribution-based Cost Inference (RCI) framework performs two steps:

1. Redistribute the trajectory-level binary signal into dense per-step costs using return decomposition.
2. Train a constrained policy (e.g., via Lagrangian dual gradient descent) on the augmented dataset with the inferred costs.

Theoretical Guarantee§

Let the original CMDP have cost function $c(s,a)$. With stop-feedback, the only observable per trajectory is: $$ \bar{C}(\tau) = \mathbf{1}[\exists t \text{ s.t. } c(s_t,a_t) > 0] $$ A redistribution $\hat{c}(s,a)$ is return-equivalent if for every trajectory $\tau$: $$ \sum_t \hat{c}(s_t,a_t) = \bar{C}(\tau) $$ The paper proves that any return-equivalent redistribution preserves:

• The feasible policy set $\Pi_C = \{\pi : \mathbb{E}_\pi[\sum_t \gamma^t c(s_t,a_t)] \le d\}$
• The optimal Lagrangian multiplier $\lambda^*$

This holds because the constraint only depends on the trajectory-level total cost, so per-step decomposition does not change the expected cost under any policy. However, the conditioning of the cost critic varies—well-chosen redistributions reduce variance and accelerate learning.

Implementation Details§

A typical implementation learns a cost-redistribution model $f_\theta(s_t,a_t,s_{t+1})$ that outputs logits. These logits are normalized via softmax within each trajectory to produce a probability distribution over steps, then scaled by the trajectory-level stop flag. This guarantees return-equivalence.

import torch
import torch.nn as nn

class CostRedistributionModel(nn.Module):
    def __init__(self, state_dim, action_dim, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim * 2 + action_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, 1)
        )

    def forward(self, states, actions, next_states):
        # states, actions, next_states: [batch, T, dim]
        x = torch.cat([states, actions, next_states], dim=-1)
        return self.net(x).squeeze(-1)  # [batch, T]

def redistribute(model, trajectories):
    # trajectories is a dict with keys: states, actions, next_states, stop_flags
    logits = model(trajectories['states'], trajectories['actions'], trajectories['next_states'])
    # softmax over time steps: per-trajectory normalized (sum to 1)
    probs = torch.softmax(logits, dim=1)
    # scale by stop flag: 1 for unsafe trajectories, 0 for safe
    costs = probs * trajectories['stop_flags'].unsqueeze(1)
    return costs.detach()

After obtaining dense costs, the training objective becomes a standard constrained RL objective: $$ \min_\pi \max_{\lambda \ge 0} \, \mathbb{E}_{\tau \sim \mathcal{D}_\text{aug}} [\sum_t \gamma^t r_t] - \lambda (\mathbb{E}_{\tau \sim \mathcal{D}_\text{aug}} [\sum_t \gamma^t \hat{c}_t] - d) $$ where $\mathcal{D}_\text{aug}$ is the original dataset with inferred costs, and $d$ is the threshold.

Results & Robustness§

Experiments on highway driving and robotic manipulation show substantially lower violation rates than:

• Sparse baselines (treat stop-feedback as a one-hot cost only at the first unsafe step)
• Classifier-based cost inference (predicting per-step cost from learned binary classifiers)

The method is robust to heterogeneous dataset compositions and label noise, suggesting practical deployment in real-world safety-critical offline RL.