Large language models (LLMs) are constantly improvising, thanks to advances in artificial intelligence and machine learning. LLMs are making significant advances in subfields of AI, including natural language processing, natural language understanding, natural language generation, and computer vision. These models are trained on massive internet-scale datasets to develop general-purpose models capable of handling a range of linguistic and visual tasks. The availability of large data sets and well-thought-out architectures that can efficiently scale with data size and models is driving this growth.
LLMs have been successfully extended to robotics in recent times. However, what remains to be achieved is a general-purpose embodied agent that learns to perform many control tasks via low-level actions from a number of large, non-curated datasets. Current approaches to general embodied agents face two major obstacles, which are as follows.
- Assumption of quasi-expert trajectories: Due to the severe limitation in the amount of available data, many existing behavior cloning methods rely on quasi-expert trajectories. This implies that agents are less flexible when faced with different tasks, because they need expert, high-quality demonstrations to learn.
- Lack of scalable continuous monitoring methods: Large, non-curated data sets cannot be effectively managed by a number of scalable continuous monitoring methods. Many existing reinforcement learning (RL) algorithms rely on task-specific hyperparameters and are optimized for single-task learning.
To address these issues, a team of researchers recently introduced TD-MPC2, an extension of the TD-MPC (Trajectory Distribution Model Predictive Control) family of model-based RL algorithms. Large, uncurated datasets spanning multiple task domains, embodiments, and action spaces were used to train TD-MPC2, a system for creating general-purpose global models. One of the important features is that it does not require hyperparameter tuning.
The main elements of the TD-MPC2 are as follows.
- Local trajectory optimization in latent space: Without the need for a decoder, TD-MPC2 performs local trajectory optimization in the latent space of a trained implicit world model.
- Algorithmic robustness: By remaking important design decisions, the algorithm becomes more resilient.
- Architecture for multiple embodiments and action spaces: Without requiring prior domain expertise, the architecture is carefully created to support datasets with multiple embodiments and action spaces.
The team shared that after evaluation, TD-MPC2 performs consistently better than model-based and model-free approaches that are currently used for various continuous monitoring tasks. This works particularly well in difficult subsets such as pick and place and locomotion tasks. The increased capabilities of the agent demonstrate its scalability as the size of models and data increases.
The team has summarized some notable features of the TD-MPC2, which are as follows.
- Improved Performance: When used on a variety of RL tasks, TD-MPC2 offers improvements over baseline algorithms.
- Consistency with a single set of hyperparameters: One of the main advantages of TD-MPC2 is its ability to reliably produce impressive results with a single set of hyperparameters. This streamlines the adjustment procedure and facilitates application to a range of jobs.
- Scalability: The capabilities of the agent increase as the size of the model and data increases. This scalability is essential to handle more complex tasks and adapt to various situations.
The team trained a single agent with a massive parameter count of 317 million to complete 80 tasks, demonstrating the scalability and efficiency of TD-MPC2. These tasks require multiple embodiments, that is, physical forms of the agent and action spaces in multiple task domains. This demonstrates the versatility and strength of the TD-MPC2 to solve a wide range of difficulties.
Tanya Malhotra is in her final year of undergraduate studies at the University of Petroleum and Energy Studies, Dehradun and is pursuing BTech in Computer Engineering with specialization in Artificial Intelligence and Machine Learning.
She is passionate about data science, with good analytical and critical thinking, as well as a keen interest in learning new skills, leading groups and managing work in an organized manner.