🤖 machine learning

RS-Diffuser: Risk-Sensitive Diffusion Planning with Distributional Value Guidance

This paper introduces RS-Diffuser, a risk-sensitive offline diffusion planning framework that integrates distributional value critics with tail-aware guidance to enable a single model to flexibly generate risk-averse, neutral, or seeking behaviors while improving robustness and reducing safety violations in safety-critical applications.

Original authors: Shiqiang Gong

Published 2026-06-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Shiqiang Gong

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to navigate a complex city. You have a massive video library of how other robots (or humans) drove through this city in the past, but you cannot let the robot drive around in the real world to learn, because one mistake could cause a crash. This is the world of Offline Reinforcement Learning: learning from a fixed dataset without new experiments.

The problem with many current AI planners is that they are "risk-neutral." They are like a GPS that only cares about the average fastest route. If a route has a 99% chance of being fast and a 1% chance of hitting a massive pothole that destroys the car, a risk-neutral planner might still choose it because the average time is good. In safety-critical situations (like self-driving cars or medical robots), that 1% chance of disaster is unacceptable.

Enter RS-Diffuser, a new method that acts like a "safety-conscious co-pilot." Here is how it works, broken down into simple concepts:

1. The "Dreamer" (The Diffusion Planner)

Think of the core of this system as a Dreamer. In the past, AI planners often tried to guess the next single move (like a chess player thinking one step ahead). RS-Diffuser is different; it "dreams" entire future scenarios at once.

It uses a technique called Diffusion, which is similar to how an artist might start with a canvas full of static noise and slowly refine it into a clear picture.

The Process: The AI starts with a chaotic, noisy guess of what the future path looks like. It then iteratively "denoises" this path, cleaning it up step-by-step until it reveals a smooth, logical trajectory of where the robot should go.
The Benefit: Because it generates the whole path at once, it understands long-term consequences better than methods that just look at the next second.

2. The "Risk Auditor" (The Distributional Value Critic)

This is the paper's secret sauce. Most AI critics (judges) only give a single score: "This path is worth 100 points."
RS-Diffuser's critic is a Risk Auditor. Instead of giving one score, it gives a full report card of possibilities.

It asks: "What is the best case? What is the worst case? What happens 90% of the time? What happens in the rare 1% of disasters?"
It uses a statistical tool called Quantile Regression to map out the entire spectrum of possible outcomes, not just the average.

3. The "Steering Wheel" (Risk-Sensitive Guidance)

Here is where the magic happens. Usually, once an AI model is trained, its personality is fixed. If you want it to be safe, you have to retrain it from scratch.

RS-Diffuser changes the rules. It separates the Dreamer (the planner) from the Auditor (the critic).

At Training Time: The model learns to dream up paths and the auditor learns to grade them based on all possible outcomes.
At Test Time (When the robot is actually moving): You can turn a "Risk Dial" without retraining anything.
- Risk-Averse Mode: You tell the system, "I care about the worst-case scenario." The system uses the Auditor's data to look at the bottom of the report card (the disasters) and steers the Dreamer away from any path that might lead there. It becomes very conservative.
- Risk-Neutral Mode: You tell it, "Just give me the average best path." It ignores the disasters and focuses on the mean.
- Risk-Seeking Mode: You tell it, "Go for the high reward, even if it's dangerous." It steers toward paths with high upside, ignoring the potential downsides.

The Analogy: The Weather Forecast

Imagine you are planning a picnic.

Old AI (Risk-Neutral): Looks at the weather forecast and says, "The average temperature is 75°F. Let's go!" It ignores the 1% chance of a tornado.
RS-Diffuser (Risk-Sensitive): The "Auditor" gives you the full forecast: "75°F average, but there is a 1% chance of a tornado and a 10% chance of heavy rain."
- If you set the dial to "Safety First," the system says, "No, the tornado risk is too high. Let's stay home."
- If you set the dial to "Adventure," it says, "The average is great, let's go!"
- Crucially, the system didn't need to learn a new way to predict weather; it just used the same prediction data but interpreted it differently based on your settings.

What the Paper Claims

The authors tested this on two main types of challenges:

Simulated Robot Control (D4RL): Robots like "Half-Cheetah" or "Walker2D" that have to move fast but might fall over or break if they move too aggressively.
Risky Navigation: Robots trying to reach a goal while avoiding "danger zones" that give huge penalties.

The Results:

Better Safety: RS-Diffuser caused fewer crashes and safety violations than previous methods.
Better Performance: It didn't just play it safe; it actually achieved higher scores (returns) than other risk-sensitive methods because it could balance safety and reward more effectively.
Flexibility: A single trained model could switch between being a cautious driver, a normal driver, or a reckless driver just by changing a number at the moment of decision, without needing to be retrained.

In short, RS-Diffuser is a planning system that doesn't just guess the future; it understands the risks of the future and lets you decide how much risk you are willing to take, all from a single, pre-trained model.

Technical Summary: RS-Diffuser

Problem Statement

Offline Reinforcement Learning (RL) aims to learn decision-making policies from fixed datasets without further environment interaction, a capability critical for safety-sensitive domains like healthcare and robotics where online exploration is costly or unsafe. However, existing offline RL methods face two primary challenges:

Distributional Shift: Policies often act outside the support of the behavior data, leading to unreliable value estimates.
Risk Neutrality: Most current diffusion-based planners and offline RL methods optimize for expected return. They fail to explicitly model the tail behavior of returns, making them susceptible to rare but catastrophic outcomes. In safety-critical settings, two behaviors with similar expected returns may possess drastically different downside risks, a distinction standard methods overlook.

While Distributional RL offers a principled way to model full return distributions (enabling metrics like Value-at-Risk and Conditional Value-at-Risk), existing risk-sensitive offline methods are predominantly built on actor-critic architectures. There is a lack of work exploring how to imbue diffusion-based planners with explicit, controllable risk sensitivity.

Methodology: RS-Diffuser

The authors propose RS-Diffuser, a risk-sensitive offline diffusion planning framework that integrates diffusion-based trajectory generation with distributional value learning. The framework consists of three core components trained on an offline dataset:

State Trajectory Diffusion Planner:
Instead of generating state-action pairs directly, the planner models a conditional distribution over future state trajectories $p_\theta(x_0 | s_t)$ , where $x_0$ represents a sequence of future states $(s_{t+1}, \dots, s_{t+H})$ . This decouples long-horizon planning from low-level action generation. The model is trained using a standard denoising objective (DDPM) to reverse noise corruption on state sequences.
Inverse Dynamics Action Decoder:
A separate deterministic decoder $f_\psi$ is trained via supervised regression to map planned state transitions back to executable actions ( $a_t = f_\psi(s_t, \hat{s}_{t+1})$ ). This allows the diffusion model to focus on feasible state evolutions while the decoder handles action recovery.
Monte Carlo Distributional Value Critic:
This is the core innovation for risk sensitivity. Unlike standard critics that predict a single expected return, this critic models the full return distribution of candidate plans.
- Training: It uses Quantile Regression with Monte Carlo supervision (rather than Bellman bootstrapping) to predict $N$ quantiles of the return distribution $Z(x)$ for a given noisy plan $x_k$ , current state $s_t$ , and diffusion step $k$ .
- Loss: The critic is optimized using the Quantile Huber loss to accurately approximate the return distribution.

Risk-Guided Inference:
The key operational feature of RS-Diffuser is its ability to control risk preferences at inference time without retraining.

A risk-sensitive objective $\rho_\alpha$ (e.g., CVaR for risk-aversion or upper-tail averages for risk-seeking) is computed from the predicted quantiles of the distributional critic.
During the diffusion denoising process, a guidance signal is injected. The gradient of the risk objective with respect to the noisy trajectory ( $g_k = \nabla_{x_k} \rho_\alpha$ ) is used to steer the generation toward trajectories that satisfy the desired risk profile.
By adjusting the risk parameter $\alpha$ or the choice of functional (CVaR vs. VaR), a single trained model can seamlessly switch between risk-averse, risk-neutral, and risk-seeking behaviors.

Key Contributions

The paper outlines three main contributions:

Framework Integration: The introduction of a diffusion planning framework for offline RL that explicitly incorporates distributional value modeling to enable tail-aware decision making.
Flexible Guidance: A demonstration of how a distributional critic can be integrated into guided diffusion sampling, providing flexible, inference-time control over risk preferences without the need for model retraining.
Empirical Superiority: Experimental evidence showing that RS-Diffuser outperforms existing state-of-the-art models on risk-sensitive benchmarks, achieving better overall returns and worst-case robustness.

Experimental Results

The authors evaluated RS-Diffuser on two benchmarks: Risk-Sensitive D4RL (MuJoCo tasks with stochastic penalties) and Risky Robot Navigation (PointMass and Ant environments with hazardous zones).

Risk-Sensitive D4RL: RS-Diffuser achieved the best overall performance across Medium, Expert, and Mixed datasets for Half-Cheetah, Walker2D, and Hopper. It outperformed both risk-neutral baselines (CQL, Diffusion-QL) and risk-sensitive baselines (OWCPG, ORAAC, CODAC, UDAC). Notably, it improved both the mean return and the Conditional Value-at-Risk (CVaR $_{0.1}$ ), indicating superior performance in both average and worst-case scenarios.
Risky Robot Navigation: In the Risky Ant and Risky PointMass tasks, RS-Diffuser achieved the highest returns while maintaining zero or near-zero safety violations. It successfully balanced safety and performance, outperforming other methods that either took hazardous shortcuts (risk-neutral) or failed to optimize returns effectively.
Ablation Studies: Experiments varying the risk parameter $\alpha$ and the risk measure (VaR vs. CVaR) confirmed that the method allows for a smooth trade-off between average return and robustness. CVaR-based guidance was found to offer the best trade-off for the default configuration.

Significance and Claims

The paper claims that RS-Diffuser addresses a critical gap in offline RL by enabling controllable risk sensitivity within the generative diffusion paradigm. Unlike previous methods that require training separate models for different risk profiles or rely on fixed risk objectives, RS-Diffuser allows a single model to adapt to varying safety requirements at inference time.

The authors position this work as a step toward flexible and reliable decision-making in safety-critical applications. They argue that combining the expressive power of diffusion models (for modeling complex, multimodal trajectories) with distributional value guidance (for tail-aware evaluation) creates a robust framework for real-world deployment where avoiding catastrophic outcomes is as important as maximizing average reward. Future work is suggested to extend this approach to higher-dimensional robotics tasks and multi-objective safety constraints.