DiffOG: Differentiable Policy Trajectory Optimization with Generalizability
Zhengtong Xu, Zichen Miao, Qiang Qiu, Zhe Zhang, Yu She
TL;DR
DiffOG addresses the challenge of generating high-quality, constraint-consistent action trajectories for visuomotor policies by integrating a differentiable trajectory optimization layer with a transformer-based projection to a learnable, SPD cost matrix ${f Q}$. The method balances fidelity to demonstrations, hard constraint satisfaction, and trajectory smoothness within a quadratic program solvable by differentiable solvers, enabling end-to-end training. Through extensive simulations and real-world experiments, DiffOG outperforms post-hoc trajectory processing and existing constrained visuomotor policies, while maintaining base policy performance and achieving real-time inference. Its generalizable design supports long-horizon, high-DOF tasks and enables zero-shot adaptation to time-varying constraints, offering a principled, interpretable approach to safe and robust robot manipulation. The work highlights future directions including reinforcement-learning fine-tuning, expanded action spaces (e.g., torques, absolute rotations), and incorporating collision-avoidance constraints for cluttered environments.
Abstract
Imitation learning-based visuomotor policies excel at manipulation tasks but often produce suboptimal action trajectories compared to model-based methods. Directly mapping camera data to actions via neural networks can result in jerky motions and difficulties in meeting critical constraints, compromising safety and robustness in real-world deployment. For tasks that require high robustness or strict adherence to constraints, ensuring trajectory quality is crucial. However, the lack of interpretability in neural networks makes it challenging to generate constraint-compliant actions in a controlled manner. This paper introduces differentiable policy trajectory optimization with generalizability (DiffOG), a learning-based trajectory optimization framework designed to enhance visuomotor policies. By leveraging the proposed differentiable formulation of trajectory optimization with transformer, DiffOG seamlessly integrates policies with a generalizable optimization layer. DiffOG refines action trajectories to be smoother and more constraint-compliant while maintaining alignment with the original demonstration distribution, thus avoiding degradation in policy performance. We evaluated DiffOG across 11 simulated tasks and 2 real-world tasks. The results demonstrate that DiffOG significantly enhances the trajectory quality of visuomotor policies while having minimal impact on policy performance, outperforming trajectory processing baselines such as greedy constraint clipping and penalty-based trajectory optimization. Furthermore, DiffOG achieves superior performance compared to existing constrained visuomotor policy. For more details, please visit the project website: https://zhengtongxu.github.io/diffog-website/.