Improving Generalization of Differentiable Simulator Policies with Sharpness-Aware Optimization

This work contributes to the ongoing discussion on the trade-off between performance and generalization in reinforcement learning, particularly in the context of sim-to-real transfer in robotics. We investigate the generalization capabilities of policies learned using differentiable simulators in contact-rich robotic scenarios. While first-order optimization achieves a higher sample efficiency, it has been empirically shown to be unstable in loco-manipulation problems. We demonstrate that, while first-order methods achieve superior performance and sample efficiency in training, they exhibit less robustness to environmental variations. To address this limitation, we propose augmenting them with sharpness-aware optimization. Our simulation results show that this approach improves the generalization of learned policies over a larger magnitude of perturbation noise.