Temporal Logic Guidance for Action-Only Diffusion Policies with World Models

Abstract: Diffusion policies enable multimodal robot behavior but offer limited ability to choose among behavior modes at inference time, even though such control is desirable in human-robot settings. Prior solutions to this lack of control have utilized Signal Temporal Logic (STL) to express human intentions and provide corresponding guidance for diffusion policy inference. However, these approaches can only guide diffusion policies that jointly generate future actions and states, increasing both complexity and runtime. We propose a novel guidance method for action-only diffusion policies that uses a separate learned world model to enable differentiable evaluation of STL robustness, with its gradient then injected into the diffusion process. This steers behavior toward constraint satisfaction without retraining, improving constraint adherence while preserving task performance. On the Can Transport task from Robomimic, our method maintains 100% task success while reducing constraint violations from over 80% for baseline methods to 4%. We also discuss extensions toward improved robustness and more complex constraints.

• The paper introduces a method that injects STL constraint gradients into diffusion steps to achieve 100% task success in robotic manipulation.
• It leverages differentiable world models to predict future trajectories, reducing constraint violations from over 80% to 4% and significantly lowering average tilt.
• The approach avoids retraining policies, offering runtime steerability and extensibility for handling complex, long-horizon temporal logic constraints.

Temporal Logic Guidance for Action-Only Diffusion Policies with World Models

Overview and Motivation

This paper introduces a method for inference-time guidance of action-only diffusion policies in robotic settings, leveraging Signal Temporal Logic (STL) constraints evaluated on rollouts produced by a separately trained world model. As modern diffusion policy architectures increasingly focus on generating sequences of actions rather than joint action-state trajectories, previous STL-based guidance methods become inapplicable or computationally prohibitive. The proposed approach addresses the gap by steering action-only diffusion policies toward constraint satisfaction at inference—which is crucial for adapting robot behavior to heterogeneous, potentially dynamic user requirements in human-robot interaction scenarios.

Methodology

The guidance mechanism operates by injecting the gradient of an STL robustness objective into the denoising steps of the diffusion process. Formally, given a diffusion policy $\pi(a_{t:t+H} | s_t)$ and an STL constraint $\phi$ over state trajectories, the method ensures that the action chunk sampled by the policy induces trajectories—predicted by a world model $f_\theta$—that satisfy $\phi$. The STL robustness metric $\rho_\phi$ is differentiable w.r.t. the state trajectories, and the world model $f_\theta$ is differentiable w.r.t. actions, allowing the guidance objective $\mathcal{J}(\phi, a_{t:t+H}, s_t)$ to be incorporated directly into the diffusion process. The approach does not require retraining the underlying policy, as constraints can be specified and enforced at runtime. Empirically, the method includes additional gradient ascent steps on the final action sequence to further optimize constraint satisfaction.

Experimental Results

Evaluation was conducted on the Can Transport task from Robomimic, where diverse human demonstrations influence policy multimodality. The STL constraint $G(R_{zz} > \cos(5^\circ))$ directs the robot to keep the can nearly upright throughout its trajectory. The method achieves 100% task success while reducing constraint violations from over 80% (baselines) to 4%, with average tilt substantially lowered (from 8.51° to 1.93°). Gradient-based guidance outperforms sample-and-rank approaches employing the same robustness measure, indicating superior mode recovery and compliance. These results demonstrate robust multi-modal behavior steering without degrading primary task performance, and confirm the practical advantage of gradient-based guidance for underrepresented modes in demonstration data.

Theoretical and Practical Implications

The separation of policy and world model facilitates compatibility with existing action-only diffusion architectures and enables interactive constraint specification. This resolves constraints on inference-time STL guidance imposed by policies lacking state trajectory predictions. The approach is extensible for more complex STL specifications and offers a foundation for supporting long-horizon temporal constraints—an established limitation in current literature (2606.22729). The possibility of leveraging the automaton abstraction of STL specifications via learned transitions, rather than environment states, invites future research into automata-level prediction for extended specification horizons.

Future Directions

Potential avenues include: enhancing the guidance process with evolutionary or second-order optimization techniques to improve stability and efficiency; expanding to more diverse manipulation tasks and constraints; and advancing toward automaton-level guidance for arbitrarily long and complex tasks. Further, exploring mechanisms for harnessing the automaton state-space for constraint evaluation could enable scalable guidance for temporal logic specifications beyond the immediate action horizon, circumventing bottlenecks of environment-level prediction.

Conclusion

The paper presents a principled and effective method for inference-time STL-guidance of action-only diffusion policies, using differentiable world models to predict future state trajectories. The method achieves strong empirical constraint satisfaction without sacrificing task success, and lays theoretical groundwork for broad, runtime-steerable policy adaptation in multimodal robotic tasks. The framework's generality and extensibility have substantive implications for safe, interpretable, and user-adaptive robotic control in variable human environments.