ConTent: Contact-Force Prediction for Latent World Models Enhances Dynamic Manipulation

Latent world models are powerful encoders of both observations and environment dynamics in robot learning. This study applies a world model to a dynamic manipulation task in which a robot uses a stick to push a ball rolling down a tilted board up into a goal zone. In this task, it is crucial that the observed task-relevant world state, such as the ball state, be accurately encoded in the latent space. Yet, although the regularisers of world models prevent collapse of the latent space, the representation can still degenerate and discard such information, thereby degrading the performance of model predictive control. We show that the discarded information can be recovered by predicting contact forces, which anchor the object representation in the latent space. We exploit this recovery to propose ConTent, contact-force prediction for latent world models, which outperforms the LeWorldModel baseline on our manipulation task. Our method predicts contact forces at training time and raises the success rate from 2% to 52% when planning with the cross-entropy method. Fine-tuning the predictor on on-policy data further improves it to 77%.