Robots that learn manipulation skills by imitation copy what they have seen, but they rarely recover when something goes wrong, and collecting enough real-world experience for reinforcement learning (RL) is slow, costly and sometimes unsafe. World models offer a way around this. A world model is a neural network that learns to predict how a scene changes when the robot acts, so that a policy can plan ahead or practice inside the model instead of on the real robot.
This project studies how action-conditioned world models and RL can combine to learn robust manipulation skills, focusing on contact-rich tasks such as insertion and assembly. We address three questions. First, how accurate does a world model need to be before policies trained inside it work well? We will train world models of increasing size on public robot datasets and on demonstrations recorded in our lab at KTH and relate their long-horizon prediction error to the success of the resulting policies. Second, how should we combine imagined experience from the world model with experience from physics simulation? We will compare model-free RL in GPU-parallel simulation, model-based RL in imagination, and hybrid schemes that use the world model for short rollouts and the simulator for correction. Third, do the gains carry over to real robots? The best policies will be tested on real robot arms in our lab.
The expected outcomes are trained world models and policies that we will release openly, a systematic comparison of model-based and model-free RL for manipulation, and conference and journal publications in robot learning and machine learning.
All training runs on GPUs. We request 5,000 GPU-h/month on Arrhenius GPU for six months, together with 2,000 GiB of project storage. We will use the resources for world-model training, RL experiments across several random seeds, and large-scale evaluation. No sensitive or personal data are processed.