World models generate improved experience for offline robot policies
Generating Critic-Guided Synthetic Experience in Latent World Models
Machine LearningArtificial Intelligence
Summary
Collecting new data for training robots or AI agents can be expensive or risky in real life. The authors work on using world models, which simulate how things change over time, to create artificial experiences that improve learning. They developed a method where a critic—a model that judges action quality—is used to pick and refine actions inside the simulated world before adding these examples to the training data. This approach boosts performance in robotic manipulation tasks without needing more real-world trials.
What this means in practice
- •For robot control engineers: Improve robot manipulation policies offline by generating refined synthetic data from learned world models without additional real robot tests.
- •For autonomous system developers: Enhance policy training for safety-critical systems using synthetic experiences refined by critics in latent dynamics models to reduce costly real-world data needs.
Authors
Jorge de Freitas, Christos Ziakas, Alessandra Russo
Abstract
Learning from offline datasets is limited in scalability, as collecting new experience is often costly or unsafe in the real world. World models offer an alternative way to generate new experience through a learned dynamics model, yet learning policies in latent space is limited due to compounding errors in imagined rollouts. To this end, we propose using the world model as a local transition generator, leveraging critic guidance to choose actions inside that model to generate synthetic experiences. We study action refinement methodologies, introducing Critic-Guided Synthetic Experience (CGSE) generation in world models. Specifically, CGSE augments an existing offline dataset by sampling candidate actions using an offline policy and refining each to maximize a first-order approximation of the $Q$-function within a trust region. We then use a frozen latent world model to imagine the future state given each updated action, and use the critic to choose the best one, appending this synthetic transition to the replay buffer. Our method requires no additional real-environment interaction or test-time planner, as it generates synthetic experience offline using the world model. On three OGBench manipulation tasks, CGSE improves an offline policy by a relative improvement of $7.1\%$ on the average success rate.