Safety-Aware Shared Autonomy via World-Model Constrained Planning

1 Introduction

Shared autonomy has emerged as a practical paradigm for human-machine collaboration, with demonstrated benefits in robotic teleoperation, drone control, and assisted driving [Javdani et al., 2015; Reddy et al., 2018; Li and others, 2023]. Its central objective is to improve task performance and reduce operator workload through timely assistance while preserving the human operator's intent and sense of control [Du et al., 2020; Myers et al., 2024].

* Corresponding authors.

In safety-aware and long-horizon tasks, insufficient assistance can lead to unsafe scenarios and poor task performance, while excessive assistance can hinder the human operator's intent [Tan et al., 2022]. Prior work can be broadly categorized into two paradigms. One paradigm relies on an intervention function to decide when to assist the human operator [Oh et al., 2025; McMahan et al., 2024]. For example, IDA estimates the expected values of the agent and human actions to decide whether to execute the agent action instead of the human action [McMahan et al., 2024]. Nevertheless, these methods typically make intervention decisions from the current observation without explicitly anticipating multi-step future risks, which can be insufficient for long-horizon and safety-aware tasks [Huang et al., 2024; Xu et al., 2021]. The other paradigm combines expert and human actions in a stepwise manner for task execution, which improves task performance through continuous expert correction [Sun et al., 2025; Yoneda et al., 2023; Schaff and Walter, 2020]. DiffusionSA learns an expert action distribution and uses a diffusion model to correct human actions toward this distribution, while preserving the user's control authority [Yoneda et al., 2023]. However, the assistance can become overly frequent, substantially reducing human autonomy. These limitations (Fig. 1(a,b)) raise a central open question in safety-aware shared autonomy: how to jointly decide intervention timing and design an action correction mechanism in long-horizon, safety-critical tasks.

A desirable assistance mechanism should anticipate multistep future risks when deciding intervention timing and provide continuous, gradual corrections once intervention is needed. World model constrained planning provides a natural and principled approach to supporting this process [Wang et al., 2025; Vasudevan et al., 2025]. In particular, a world model enables decision making through imagined rollouts, allowing the assistant to reason about future outcomes along candidate trajectories.

Building on these insights, we propose World Model Assisted Safety Planning (WASP), a world model constrained planning framework for safety-aware human-machine collaboration. We first pretrain a safety-aware world model that consists of an RSSM-style latent dynamics model with a reward head and a collision cost head [Huang et al., 2024].

WASP is structured into two stages, Model Predictive Intervention (MPI) and Model Predictive Residual Planning (MPRP). From the perspective of intervention timing, a safety-aware intervention function can avoid collisions, but it often leads to conservative braking behavior that degrades task performance (Fig. 1(c)). Therefore, we design both safety-aware and task-aware intervention functions to address this trade-off in MPI (Fig. 1(d)). The safety-aware intervention function uses imagined rollouts in the pretrained world model to estimate cumulative collision cost and triggers intervention when the cumulative cost is sufficiently high. The task-aware intervention function uses a limited number of imagined rollouts to estimate the long-horizon return gain of a short correction and intervenes only when the expected gain exceeds a threshold. When either function triggers, WASP intervenes and transitions to MPRP.

In MPRP, WASP samples residual action sequences and evaluates their imagined trajectories in the world model, retaining candidates that satisfy a predicted cumulative collision-cost constraint. If enough safe candidates remain, WASP optimizes task return while penalizing deviation from the operator command. Otherwise, it prioritizes reducing predicted collision cost.

We evaluate WASP on Safety Gymnasium [Ji et al., 2023], comparing it against value-function-based intervention methods and stepwise action-correction methods. Across three tasks, WASP achieves strong task performance under safety constraints while requiring fewer and smaller interventions.

2 RelatedWork

Shared autonomy. Shared autonomy combines human inputs with agent assistance to improve task performance while preserving human intent [Abbink et al., 2018]. Prior work includes intent inference, empowerment-based assistance, action correction, and intervention burden [Weiss and Gerdes, 2023]. Among these methods, action correction and intervention quantity have received particular attention. One line of work studies action correction by overriding or adjusting user actions to improve operator task performance [Schaff and Walter, 2020; Li and others, 2023]. Another line explicitly models intervention quantity, e.g., optimizing assistance under hard or soft constraints on the number of interventions to improve user experience and reduce manual tuning [Tan et al., 2022; Broad et al., 2019]. However, safety considerations are often treated as secondary in much of this literature.

Safety-aware shared autonomy. Safety-aware shared autonomy has mainly been explored through two paradigms: action correction and intervention-function-based assistance [Oh et al., 2025; McMahan et al., 2024]. Action-correction methods adjust human actions toward an expert action distribution. For example, diffusion-based approaches map human actions to samples from this distribution [Yoneda et al., 2023; Sun et al., 2025], while residual policy learning trains a model-free assistive policy to minimally adjust human actions under goal-agnostic constraints [Schaff and Walter, 2020; Li and others, 2023]. These methods can improve task performance, but their continuous corrections may lead to overly frequent assistance and reduced operator autonomy [Tan et al., 2022; Broad et al., 2019]. In contrast, interventionfunction-based methods decide when the agent should override human actions using value or safety-margin signals, as in interventional assistance [McMahan et al., 2024; Xue and others, 2023; Peng and others, 2023] and human-centered safety filters [Oh et al., 2025]. However, these decisions are often based on current-state value or margin estimates, which may be insufficient when hazards are delayed and require multi-step anticipation. Moreover, some safety-filter approaches require expert priors to construct safety margins [Oh et al., 2025].

Safety-constrained planning with latent world models. Safety-constrained planning with latent world models offers a natural way to anticipate delayed hazards under partial observability [As et al., 2025; Huang et al., 2024]. Latent dynamics models support imagined rollouts for evaluating future outcomes over candidate action sequences. Such rollouts have been widely used for imagination-based control and model-predictive planning [Wang et al., 2025; Vasudevan et al., 2025]. Recent safety-oriented model-based work augments this pipeline with cost or risk prediction and uses predicted cumulative cost to constrain planning or filter unsafe rollouts [Zwane et al., 2023; Hogewind et al., 2022; Jayant and Bhatnagar, 2022].

Relation to prior work. Prior shared-autonomy methods typically address either action correction or value-/marginbased intervention, while safety-constrained latent world models enable multi-step risk-aware planning under partial observability. WASP bridges these lines by using a pretrained safety-aware world model to jointly decide when to intervene and how to correct operator commands.

3 ProblemFormulation

Constrained Markov decision process (CMDP). Safetyconstrained long-horizon decision making is commonly modeled as a constrained Markov decision process (CMDP), defined as M = (S,A, p, r, c, γ, κ) [Chow et al., 2018; Dai et al., 2023]. The state and action spaces are S and A, respectively. The transition kernel p(s′ | s, a) gives the probability of reaching s′ ∈ S after taking a ∈ A in s ∈ S. The reward function r(s, a) specifies the immediate reward, and the cost function c(s, a) specifies the immediate safety cost. The discount factor is γ ∈ (0, 1) and κ denotes the allowable discounted cost budget. The discounted return and cumulative cost of a policy π(a | s) are

The CMDP objective is

Safety-aware shared autonomy. Based on the CMDP formulation, we formulate safety-aware shared autonomy by coupling a nominal human command with an assistive residual correction [Schaff and Walter, 2020]. At time t, the operator provides a nominal command aHt ∈ A and the assistant outputs a correction δat. The executed action is at = aHt + δat, followed by projection or clipping to the valid action bounds when necessary.

To make the shared-autonomy information structure explicit, we augment the state with the operator command,

and define the assistive policy as a residual policy conditioned on ˜ st:

The environment evolves according to the same transition kernel p(st+1 | st, at) under the executed action at.

We implement minimal intervention by penalizing correction magnitude while enforcing the cumulative safety-cost constraint. This yields the objective

where α > 0 controls the penalty on correction magnitude and κ specifies the allowable discounted safety-cost budget.

4 Methods

Our method combines a pretrained safety-aware world model with online intervention gating and residual planning in a closed loop. Figure 2 provides an overview of the entire framework, from environment observation to the execution of the final corrected action, and the subsequent update of the state for the next time step.

4.1 Pretrained Safety-AwareWorldModel

To pre-evaluate future outcomes for intervention timing and action correction, we pretrain a safety-aware world model offline, together with reward and cost critics ˆ Vr and ˆ Vc, and keep them fixed during online shared-autonomy control.

The world model is parameterized as a recurrent statespace model (RSSM), which summarizes history with a latent state consisting of a deterministic recurrent component ht and a stochastic latent variable zt. Given (ht, zt) and the executed action at, the model updates the recurrent state via ht+1 = fθ(ht, zt, at) and predicts the next stochastic latent state through the prior distribution pθ(zt+1 | ht+1). When an observation ot is available, the stochastic latent variable is inferred through the posterior distribution qθ(zt | ht, ot). During imagination, the model rolls forward using only the prior dynamics.

From the latent state (ht, zt), the world model reconstructs the observation via a decoder Dθ(ot | ht, zt) and predicts immediate reward and safety cost through decoders Rθ(rt | ht, zt) and Cθ(ct | ht, zt). We denote the predicted means by ˆ rt and ˆ ct, where ˆ ct ∈ [0, 1] represents the predicted probability of a safety-cost event. Two value critics, ˆ V ψr r (ht, zt) and ˆ V ψc c (ht, zt), provide bootstrapped estimates of long-horizon reward and cost beyond the finite imagination horizon.

The model is pretrained on replayed sequences (ot, at, rt, ct) by minimizing

where sg(·) denotes the stop-gradient operator. The first two terms correspond to the representation regularization and future-prediction losses of the RSSM, encouraging consistency between the posterior and the learned prior. The remaining terms are negative log-likelihoods for observation reconstruction, reward prediction, and cost prediction, weighted by βo, βr, and βc, respectively.

The reward and cost critics are trained separately using bootstrapped TD-λ return targets computed from latent rollouts.

4.2 Model-Predictive Intervention Function

Intervention is determined by two model-predictive intervention functions. The safety-aware intervention function has priority and gates intervention based on whether nominal operator behavior is predicted to exceed the safety-cost budget. When safety is not at risk, the task-aware intervention function gates intervention based on whether a small correction is expected to yield a meaningful long-horizon return gain. This mechanism avoids unnecessary assistance while preventing overly conservative behavior that preserves safety at the expense of task progress.

At each time step, the assistant decides whether to intervene on the operator command, prioritizing safety while remaining passive when intervention is unnecessary. This decision is made by simulating nominal operator behavior through imagined rollouts in the pretrained world model.

To roll out nominal behavior, we follow the standard shared-autonomy evaluation setup and use a simulated nominal operator policy πH to generate future operator commands during imagination. At each imagined step, the future nominal command is sampled from πH conditioned on the imagined latent state, aHt+k ∼ πH(· | h(m) t+k , z(m) t+k ). In our experiments, πH is constructed from an RL-trained expert policy with injected imperfections, including action noise, lagged execution, limited skill, and delayed response. The same πH is used consistently across imagined rollouts and realenvironment evaluation for a given operator setting.

Starting from the current filtered latent state (ht, zt), we sample M imagined trajectories of horizon H under πH and the world model dynamics, producing per-step predictions {ˆ r(m) t+k }H-1 k=0 and {ˆ c(m) t+k }H-1 k=0 . Each rollout m is summarized by TD-λ returns with critic bootstrapping. For a generic signal x ∈ {r, c} with per-step prediction ˆx(m) t+k and critic ˆ Vx, initialize

and for k = H - 1, . . . , 0 compute

Here, λ ∈ (0, 1] controls the rollout-bootstrap trade-off. Setting x = r yields a model-predictive estimate of long-horizon task return, and setting x = c yields a model-predictive estimate of discounted cumulative safety cost.

Safety-aware intervention function. Safety intervention takes priority and is driven by the CMDP constraint on discounted cumulative cost. Under nominal operator behavior, we compute the mean TD-λ cost return across M imagined rollouts:

Intervention is activated whenever ¯G(c) t,λ > τc, where τc is the model-predictive cost threshold. When this condition holds, the task-aware gate is skipped.

Task-aware intervention function. When the safetyaware intervention function is inactive, task intervention is applied only if a short residual correction is expected to yield a meaningful improvement in long-horizon return.

Let ¯G(r) t,λ denote the mean imagined TD-λ reward return under nominal operator behavior:

To test whether task intervention is warranted, we run a budgeted sampling probe over an L-step residual sequence δat:t+L-1. This probe is used only for gating and evaluates a finite set of candidate residual sequences without iterative distribution updates. Each candidate sequence is evaluated by composing it with nominal operator commands: the first L actions are corrected, and the remaining actions follow the nominal commands generated by πH in imagination. Let ¯G(r) t,λ (δat:t+L-1) and ¯G(c) t,λ (δat:t+L-1) denote the mean imagined TD-λ reward and cost returns under this composed rollout. We compute a candidate correction by solving

Denote ¯Gcorr t = ¯G(r) t,λ (δacorr t:t+L-1 ) and define

The task-aware intervention function activates only when ∆Gt > τr, where τr is the return-gain threshold.

4.3 Model-Predictive Residual Planning

When intervention is applied, residual corrections are planned online in a receding-horizon manner using imagined rollouts from the pretrained world model. At time t, we optimize an L-step residual sequence while evaluating task return and safety over a longer prediction horizon H. In practice, near the safety boundary, many sampled residual sequences lead to predicted safety-constraint violations, leaving too few safe candidates for reliable CEM distribution updates. We therefore handle planning in two cases: one where the sampled batch contains enough safe candidates, and one where safe candidates are scarce.

We plan over an L-step residual sequence δat:t+L-1 and evaluate it by composing the residuals with nominal operator commands:

where aHt is the current operator command and future nominal commands aHt+k for k ≥ 1 are generated by πH during imagination. In implementation, the composed actions are clipped to the valid action bounds before being passed to the world model.

For each sampled candidate δa(i) t:t+L-1 , we roll out the world model for horizon H and compute its predicted discounted cumulative safety cost via the mean TD-λ cost return ¯G(c),(i) t,λ , averaged over M imagined rollouts. A candidate is treated as safe if its predicted cost satisfies the modelpredictive cost threshold:

We denote the number of safe candidates by |Fs|. Planning is handled in two cases depending on whether |Fs| is at least a threshold Nf .

Case I (|Fs| ≥ Nf ). When the batch contains enough safe candidates, we update the CEM sampling distribution using elites selected from Fs. For each candidate δa(i) t:t+L-1 , we estimate its mean imagined TD-λ reward return:

where each rollout uses the composed action sequence in Eq. (13). We score safe candidates by

and update the sampling distribution using the highestscoring candidates within Fs as elites.

Case II (|Fs| < Nf ). When safe candidates are scarce, Fs provides too few elites for a reliable CEM distribution update. We therefore perform a recovery update by selecting elites that minimize predicted safety-cost violation while still penalizing large corrections.

Using the predicted cost return ¯G(c),(i) t,λ for each candidate i, we define

where [x]+ = max(x, 0). We select the Nf candidates with the smallest S(i) safe as elites to update the sampling distribution. In subsequent CEM iterations, if the sampled batch contains at least Nf safe candidates, planning switches back to Case I.

Given the selected residual sequence, the action executed in the environment is obtained by composing the nominal operator command with the first residual and enforcing action bounds:

where δa∗t is the first element of the selected residual sequence: the highest-scoring feasible candidate under Stask if feasible candidates are available, and otherwise the candidate with the smallest Ssafe. We then apply at in the real environment, observe the next transition, update the filtered latent state using the new observation, and repeat the interventionfunction gating and residual planning at time t+1. This receding-horizon loop closes the control pipeline by continuously replanning corrections from updated latent state estimates.

5 Experiments

Our experiments evaluate WASP in terms of task performance, safety, autonomy, and intervention necessity. Unless otherwise specified, results are averaged over three random seeds, and main results are aggregated over the imperfect operator variants.

5.1 Experimental Setup

We evaluate WASP on three Safety-Gymnasium tasks [Ji et al., 2023] with visual observations and explicit safety costs: SafetyCarRace1-v0, SafetyCarFormulaOne1-v0, and SafetyPointGoal1-v0. All tasks adopt distance-based reward shaping with an additional goal-completion bonus, provide front- and rear-view camera observations, and use a continuous two-dimensional control space for steering and longitudinal motion. These tasks go beyond obstacle avoidance by combining goal-directed navigation with multiple safety-cost sources, including hazards, collisions, and barrier contacts.

We use H = 15, L = 5, Nprobe = 300 for the taskaware probe without iterative updates, and CEM with population N = 300, elites Ne = 100, and I = 6 iterations for MPRP.

We compare our approach against several baselines that represent key design axes in prior work. The Humanonly baseline executes the operator's actions without assistance, serving as a reference for evaluating the effectiveness of shared autonomy. IDA (Intervention-Driven Assistance) leverages a learned Q-function to decide when to override the operator's action, substituting an expert action when necessary [McMahan et al., 2024]. Diffusion continuously applies corrections toward expert-like behavior, while SafeRLRes employs a model-free residual controller trained under safety constraints [Yoneda et al., 2023; Schaff and Walter, 2020].

To model realistic operator imperfections, we simulate four variants: noisy operators introduce stochastic perturbations, laggy operators repeat the previous action with probability plag, limited-skill operators use a less skilled expert checkpoint trained with fewer interactions, and reaction-time operators apply delayed commands. These variants capture noise, temporal inconsistency, limited skill, and sluggish response, respectively, and assess the ability of predictive intervention to prevent violations caused by sequences of suboptimal actions. Because several variants are stochastic, such as noisy actions and probabilistic lag, the nominal operator model used in planning is an approximate predictor rather than an exact copy of the realized operator behavior.

5.2 Main Results

Table 1 presents a comparison across three tasks: Race1, FormulaOne1, and Goal1, evaluated under the protocol described in Sec. 5.1. The table reports results aggregated over four imperfect operator variants: noisy, laggy, limitedskill, and reaction-time operators. We compare against several baselines: Human-only, IDA, Diffusion action correction, and SafeRL-Res. These baselines represent key design approaches such as value-based gating (IDA), stepwise correction (Diffusion), and model-free residual assistance (SafeRL-Res).

We evaluate the methods on three aspects: task performance, safety, and autonomy. Task performance is measured by episodic return and success rate, while safety is assessed by cumulative cost and violation rate. Autonomy is evaluated by the assistant's deviation from the nominal operator command [Oh et al., 2025]. Specifically, we define the induced residual as δat = at - aHt , where the intervention rate is the fraction of steps where ∥δat∥ > 0, and the correction magnitude is the empirical mean of ∥at - aHt ∥.

Across all three environments, WASP demonstrates a strong balance between safety and task performance. The effect is most pronounced in Race1 and FormulaOne1, where violations often emerge after a sequence of small, suboptimal actions. Compared with IDA, WASP substantially reduces both cost and violation rate while requiring fewer or smaller deviations from the operator command. This result highlights the benefit of predictive rollouts, which allow WASP to intervene more selectively before safety violations occur.

Compared with Diffusion action correction and SafeRLRes, WASP achieves comparable or superior task performance while substantially reducing intervention rate and correction magnitude. This indicates that short-horizon residual planning avoids the continual micro-corrections often seen in stepwise correction methods while still maintaining safety. Diffusion often requires large corrections because it pushes the current action toward expert-like behavior, which can disregard the operator's current intent.

5.3 Counterfactual Evaluation of Intervention Necessity

We assess whether interventions occur before risky events rather than being applied as persistent corrections. To do this, we perform counterfactual rollouts in the real environment at timesteps where intervention occurs. At each such timestep, we record the environment state, including relevant randomization factors and the operator's internal state. From this recorded state, we run two branches for W steps: at the first step, the assisted branch executes the method's action, while the counterfactual branch executes the nominal operator command. Both branches then continue with operator-only control for the remaining W - 1 steps.

We measure counterfactual risk using the mean cumulative cost over W steps:

We also compute the local benefit of intervention by comparing cost and reward differences between the assisted and counterfactual branches:

An intervention is counted as unnecessary if it does not reduce cost and its return improvement is below a threshold τR. We define the unnecessary intervention rate as

where Tint denotes the set of intervention timesteps.

Figure 3 shows results on Race1 with W = 15. WASP selectively intervenes in high-risk states, resulting in a low unnecessary intervention rate. In contrast, Diffusion and SafeRL-Res apply corrections even in low-risk states, leading to higher unnecessary intervention rates. IDA focuses more on high-risk states than stepwise correction methods but still intervenes in some low-risk cases. These results indicate that WASP intervenes more selectively, reducing unnecessary corrections while maintaining safety.

5.4 Ablation Study

Table 2 summarizes the ablation results, showing the impact of different components on task performance, safety, and autonomy. Removing the safety-aware gate and relying only on the task-aware gate causes a significant increase in violations, as the system may fail to intervene when the operator takes unsafe actions that still appear beneficial for task progress. This emphasizes the importance of multi-step safety screening for reducing safety violations. Conversely, using only the safety-aware gate largely preserves safety but leads to overly conservative interventions, which can hinder task progress and reduce performance.

Ablating Case II recovery degrades performance near the constraint boundary, where too few feasible candidates are available for stable CEM updates. This leads to higher safety costs and less stable corrections, especially under temporally inconsistent operator behavior, highlighting the importance of recovery updates for maintaining stability under constraints.

Disabling imagined rollouts and relying on a single-step value estimate increases intervention rates while worsening both cumulative cost and task return. This suggests that single-step estimates can trigger unnecessary interventions while missing delayed risks, leading to poorer intervention timing. These results highlight the importance of multi-step rollouts for accurate intervention timing and effective safetyaware residual planning.

6 Conclusion

We introduced World Model Assisted Safety Planning (WASP), a safety-aware shared autonomy framework that uses a pretrained world model for predictive intervention and residual planning. WASP anticipates future safety risks and applies low-deviation corrections only when needed. Experiments on vision-based benchmarks show that WASP substantially reduces safety violations and unnecessary interventions while maintaining competitive task performance. Although evaluated on driving and navigation tasks, WASP only assumes a nominal operator command, residual correction, and cumulative safety-cost constraint, making it applicable to broader shared-autonomy settings.

Acknowledgements

This work was supported in part by the NSFC under grant No. U23A20339, No. 62125305, No. 62573339, No. 62503380, Fundamental and Interdisciplinary Disciplines Breakthrough Plan of the Ministry of Education of China under grant No. JYB2025XDXM210, and Natural Science Foundation of Shaanxi Province under Grant 2025SYSSYSZD-083, the National Key R&D Program of China under Grant No. 2024YFB4707600, State Grid Corporation of China Science and Technology Projects 52060025005B-439-ZN, and Shaanxi Provincial Science Fund for Distinguished Young Scholars under Grant 2025JC-JCQN-077.

Contribution Statement

Zechen Shi, Jiawei Ye, and Zeyang Liu contributed equally to this work.