Enhancing World Models with Specialized Prediction Networks for Reinforcement Learning

Abstract

Training robots in the real-world using reinforcement learning is both expensive and risky. World Models—a simulated environment that mirrors real-world conditions—have been proved to offer an alternative to real-world training. Such simulation-based training not only reduces costs significantly but also reduces the dependency from real-world testing.

While previous studies focus on single-network architectures that predict state, reward, and episode termination as a single output, this research proposes a different approach by creating a structure based on specialized prediction networks for each of the aforementioned elements.

During the experiment, several simulated environments were used. The main results obtained showed that the specialized-network World Models were capable of learning the environment’s dynamics adequately, and that the proposed architecture outperformed single-network configurations by more effectively capturing these dynamics.

Finally, future directions are included on possible ways to enhance World Models efficiency.