What are "World Models"? Their origins and potential as the key to the future of AI

This article provides a general overview of "World Models," which hold the key to future AI technology development, while referencing the paper from which they originated and discussing their future possibilities and challenges.

About "World Models"

What are "World Models"? To answer this question briefly, they are a "framework for acquiring a model of the environment surrounding an agent through learning from observations. They are one of the technologies that "allow AI to efficiently learn and predict environmental changes and the results of actions," and can also be described as a method for realizing "imagination and subsequent learning and judgment" by AI.In other words, it is a technology that gives AI "imagination." It exists at the intersection of various fields such as generative models, reinforcement learning, deep learning, robotics, simulation, the metaverse, and cognitive psychology, and synergy with multimodal LLMs is also expected.

World Models originated from a 2018 paper titled "World Models," published by David Ha of Google Brain (at the time; currently CEO of Sakana AI) and Jürgen Schmidhuber, the creator of LSTM. This paper explores a method for simulated cars to learn autonomous driving on a racetrack. It attracted significant attention because it presented an innovative way for AI agents to model how they perceive and interact with their environment.

An agent refers to a "system that makes decisions and takes actions within an environment" in reinforcement learning. The paper proposed giving the agent an internal mechanism to interact with the environment in a predictive manner. This demonstrated that before taking an action, an agent can make a judgment and select an action after predicting, "If I take this action, this result will occur."

Proposed method in the "World Models" paper

I will briefly describe the methods used in the "World Models" paper.
The agent's internal architecture consists of three parts: Vision (V), Memory (M), and Controller (C).

A Variational Autoencoder (VAE) is used for dimensionality reduction of image frames in Vision. The VAE converts the massive amount of data the agent encounters (such as the vast amount of image data when a car is driving on a road) into a more manageable compressed format called a Latent Space. This latent space is important as a foundation for predicting future states in the Memory section. Another key point is that, unlike standard autoencoders, introducing randomness into the decoding process with a VAE makes it a model capable of generating new data variations.

In Memory, that latent space is used to predict the future state of the environment. Here, a Recurrent Neural Network (RNN) is used, and based on the compressed data (= encoded experience), it models how the latent space changes over time while taking into account the agent's current understanding and actions, using Autoregression. In other words, the agent simulates a potential future according to a time series, which is one of the innovative aspects of the proposed method.

In the Controller, actions are determined based on current data from Vision and the future state predicted by Memory. Here, the model parameters in the Controller section are updated using reinforcement learning methods. This allows the agent to determine the optimal action to obtain a rational result (in the context of reinforcement learning, to maximize rewards).

With this architecture consisting of the three parts—Vision, Memory, and Controller—the AI agent has succeeded in demonstrating high learning capabilities, such as achieving skillful driving on a racetrack in a simulated environment.

Similarities between "World Models" and "Mental Models" in cognitive psychology

Among the processes in the architecture consisting of the three parts—Vision, Memory, and Controller—what is particularly interesting is the structure in which the Controller makes decisions while learning in real-time by combining "current data" from Vision with "predictions (future data)" from Memory, which is a learned model. Although based on external information, this decision-making process is directly taking place within a so-called simulation environment created by the internal model of the Memory section. The term "World Models" is used to refer to this innovative construction of a simulation environment and the decision-making using it.

If we compare this to humans, it is similar to the process of imagining in one's head based on the state in front of one's eyes and previously learned experiences and knowledge, thinking about the actions to take based on the results of that imagination, and making a decision. This corresponds to what is called a "mental model" in cognitive psychology.

In short, "World Models" can also be rephrased as a method for AI to acquire its own "mental model."

The "imagination and subsequent learning" known as "generative understanding" in World Models

In understanding World Models, it is important to note that the AI agent does not simply have a "world (simulation environment)" inside it. The point is that it incorporates a process that can be called "generative understanding," just as humans can learn by exercising their imagination in a "mental model."

For example, when a person thinks about shooting a basketball, they create the execution of an imaginary shot and its results in their head (while trying and failing many times) based on physical cues, including their own vision. Even if they have never shot a basketball before, they can use their imagination to generate these images. By doing this (exercising imagination in a "mental model"), people can deepen their understanding of shooting even without actually taking a shot.

Although it is an old example, the book "Basketball FundaMENTALs: A Complete Mental Training Guide Paperback," published in 1991 (a book that explains the importance of training with mental models in basketball), cites a study at Clark University. According to it, there was a report that the same increase in shooting success rate was seen in a comparison between college students who practiced shooting for 20 minutes a day for 20 days and college students who practiced shooting within a mental model (i.e., in their "imagination").

This progress in understanding through imagination, which we will call "generative understanding" here because it involves generating fictional images to deepen understanding, is similar in that World Models deepen understanding by predicting the results of an agent's actions based on generation in decoding from past data. This allows for higher learning efficiency than general methods.

Learning by interacting with a simulation environment that incorporates the process of generative understanding in World Models is very beneficial from the perspective of learning efficiency. This process makes it possible to practice and refine skills in a safe, controlled virtual space before actually trying those actions. While that itself is the same as learning with a human "mental model," in the case of AI, it can perform learning using a virtual internal model much more efficiently and in much larger quantities than by actually experiencing and learning in the real world. This is because learning in a simulation environment is far more time-efficient than actually driving an autonomous vehicle to learn, and it becomes possible to accumulate tens of thousands of driving experiences. Therefore, a greater effect on accuracy improvement can be expected.

Updating World Models and Improving Decision-Making Capabilities

By running a car tens of thousands of times within a virtual simulation environment, the system gains vast experience. It can collect massive amounts of driving and accident data, and just as one might increase the resolution of an image, it is (theoretically) possible to update the world model into a more sophisticated one. This update would allow the agent to appropriately estimate the situation and potential future events even when it has limited information. An AI that has acquired a high-level understanding of the "world in which a car drives" will be able to imagine what might happen next, even in rare situations, and skillfully execute decisions on what actions to take.

The Versatility of World Models and the Path Ahead

Looking further into the internal architecture of the agent proposed in the "World Models" paper, one realizes its versatility. While the world model is constructed internally within the Memory section, this world model itself is independent of any specific task and is universally applicable. In other words, at the initial stage, the AI can generally learn about the environment and its actions, focusing on understanding the world it inhabits and its dynamics. Then, once an understanding of the world is established, it is possible to guide the AI's behavior to optimize for a specific task or set of values (reward function) by changing the reward settings in the Controller section. Moreover, because the world model exists, this learning becomes efficient. This separation of the world and the task is a key factor that will enhance the versatility of AI, enable the execution of more diverse and complex tasks, and ultimately lead to the development of AGI.

World models are not limited to physical control and decision-making in real space, such as driving a car or shooting a basketball. For example, they can be applied to intellectual and creative activities, such as imagining strategies for solving mathematical equations or thinking about how to proceed with a drawing to make it more creative. World models enable planning and execution based on an understanding of the world specific to each task.

Yann LeCun, a pioneer of CNNs (Convolutional Neural Networks) that revolutionized image recognition AI, the creator of LeNet, and an authority on modern AI, advocates for the importance of world models in the quest to artificially acquire the intelligence possessed by animals and humans. He is working on building a complete cognitive architecture based on world models. The following paper, "A Path Towards Autonomous Machine Intelligence," organizes his thinking and framework, demonstrating the further generalization of world models.

Reference) Paper"A Path Towards Autonomous Machine Intelligence" and related materials

Regarding OpenAI's Video Generation AI, Sora

Changing the subject slightly, on February 15, 2024 (local time), OpenAI announced "Sora," an AI model that generates videos up to one minute long from text prompts. Because it can create realistic and expressive videos, it was reported by various media outlets in Japan the following day, including TV and newspapers.

In initial reports, it was often described as a "video generation AI," as seen in this news footage, but in fact, OpenAI itself does not describe Sora that way, stating that it is being developed as a "world simulator."On the webpage introducing Sora, it is explained as a world simulator equipped with the ability to simulate humans, animals, and environments in the physical world, as well as actions that can affect the state of the world.

In other words, at OpenAI, Sora is recognized and developed as a technology for building world models. They also state, "We believe that Sora will be an important milestone in achieving AGI." This is an area of interest, including their technology development roadmap and their approach toward that goal.

Synergy Between World Models and LLMs

While world models are a promising technology, they are not capable of making effective decisions in every conceivable problem. Even for problems that are not particularly difficult for humans, there are problems that are difficult to solve using only the methods proposed by David Ha and his colleagues (problems similar to Moravec's paradox). (Yann LeCun envisions building a complete architecture that overcomes such problems, but here I will discuss the synergy with LLMs.)

The method devised by David Ha and his colleagues relies on generative understanding and adaptation through reinforcement learning, but for areas where these are difficult to solve, there is potential for synergy by combining them with the recognition and evaluation capabilities of LLMs (especially multimodal LLMs). For example, the following Hatena Blog entry cites the Atari game "Montezuma's Revenge," which is very difficult to solve with reinforcement learning-based approaches (but for which humans can easily see a path to a solution), and shows an example of smoothly solving the problem by providing direction through environmental interpretation by a multimodal LLM. In other words, this is a method of using an LLM to determine the general direction of a strategy and then using reinforcement learning to solve it.

However, currently, LLMs have the issue of potentially slow processing speeds, and there are cases where the synergy between LLMs and world models cannot be applied in situations requiring real-time processing. This can be expected to improve with compact LLMs or on-premise LLMs, but even then, there may be applications that are not effective depending on the type of problem. The following paper points out the issue of bias in LLM training data in decision-making for autonomous driving, suggesting that it is premature to have excessive expectations for LLMs.

In short, the key is how to select use cases where a world model alone is difficult to solve but can be effectively solved by incorporating LLM direction, but as a means to expand the potential of world models, the significance of combining them with LLMs is substantial.

World Models and Improving Reliability

While world models have high potential, discussions on how to build safe, transparent, and reliable architectures in the future will also be important.

For example, in autonomous vehicles applying world models, if transparency and interpretability are insufficient, there is the problem that it becomes difficult to identify the cause of a phenomenon and eliminate it in the event of a failure. Furthermore, when attempting to use world models to advance corporate decision-making systems or diagnostic systems, it may be necessary to demonstrate that there are no biases that could be problematic within a social context.

Research results on building so-called Explainable AI (XAI) should also be reflected in world models. In addition, human wisdom is required in determining what values (reward functions) should be used to guide the optimization of AI behavior. Building a mechanism to continuously reflect human insights into world models is also essential. Integration with the know-how cultivated in Human-in-the-loop learning will also be indispensable.

Conclusion

In this article, we have provided an overview of World Models, their characteristics, and their potential. World Models represent a groundbreaking approach supported by a comprehensive architecture that equips AI with the ability to imagine and, through that, to learn and make judgments. This provides a foundation for giving AI versatility and the ability to solve diverse and complex tasks.

World Models can be considered one of the technologies that will determine the future direction of AI development. Furthermore, they are expected to contribute to the understanding and analysis of the "Theory of Mind," a concept discussed in developmental psychology and primate research that explains how humans, during their growth process, acquire the recognition that others also possess a "mind" and that it is different from their own. They are also similar to Predictive Processing (PP), a theory proposed by researchers such as Andy Clark as part of modern consciousness research, suggesting that this approach may also bring us closer to understanding the problems of human consciousness and the mind.

The nature of a World Model is determined by how we internalize the world within AI and how we implement "generative understanding," which serves the function of "imagination." Depending on the implementation, it may be possible to significantly raise the baseline performance of AI beyond current levels. There is no doubt that the advancement of World Model research will pave the way for the development and deployment of advanced AI applications in various fields.

That said, while Yann LeCun states that World Models are an important theme for AI development over the next decade or so, he also notes that the limitations of World Models (or the World Model-based cognitive architecture he proposes) are not yet well understood. At the same time, because of the high level of learning and predictive performance that can be achieved, how to build World Models in a safe and fair manner becomes a critical issue. Only by safely realizing World Models will we be able to evolve AI into a reliable next-generation partner that supports society.

Related Articles