This paper presents a novel approach for optimal control of nonlinear stochastic systems using infinitesimal generator learning within infinite-dimensional reproducing kernel Hilbert spaces. Our learning framework leverages data samples of system dynamics and stage cost functions, with only control penalties and constraints provided. The proposed method directly learns the diffusion operator of a controlled Fokker-Planck-Kolmogorov equation in an infinite-dimensional hypothesis space. This operator models the continuous-time evolution of the probability measure of the control system's state. We demonstrate that this approach seamlessly integrates with modern convex operator-theoretic Hamilton-Jacobi-Bellman recursions, enabling a data-driven solution to the optimal control problem. Furthermore, our statistical learning framework includes nonparametric estimators for uncontrolled forward infinitesimal generators as a special case. Numerical experiments, ranging from synthetic differential equations to simulated robotic systems, showcase the advantages of our approach compared to both modern data-driven and classical nonlinear programming methods for optimal control.

Simulators for mechanical systems are widely used, for example in testing and verification [1, 2], model-based control strategies [3, 4], or learning-based methods [5, 6]. However, many common simulators have numerical difficulties with more complex mechanical systems involving constraints [7]. Such constraints can represent joints connecting rigid bodies, which may form kinematic loops, for example, in exoskeletons. Constraints can also be used to confine the movement of bodies, for example, to model joint limits in robotic arms, or to describe contact with other bodies or the environment in walking and grasping. Exactly enforcing such constraints can cause numerical issues, for example, due to the stiff nature of contact interactions. To alleviate these numerical issues, simulators often allow small constraint violations by representing all constraints as spring-damper elements as in MuJoCo [8] and Brax [9], or by accepting interpenetration of bodies as in Drake [10] and Bullet [11]. Small violations can sometimes be acceptable, for example, contact interpenetration in the order of micrometers for meter-scale walking robots. But millimeter or centimeter violations, for example in MuJoCo, can be considered too large. Employing these methods and accepting larger constraint violations for stable simulations contributes to the sim-to-real gap, a major issue in robotics [12].

Dojo achieves stable simulation at low sample rates and conserves energy and momentum by employing a variational integrator. A nonlinear complementarity problem with second-order cones for friction models hard contact, and is reliably solved using a custom primal-dual interior-point method. Special properties of the interior-point method are exploited using implicit differentiation to efficiently compute smooth gradients that provide useful information through contact events. We demonstrate Dojo with a number of examples including: planning, policy optimization, and system identification, that demonstrate the engine's unique ability to simulate hard contact while providing smooth, analytic gradients.