DDMRL Project

View on GitHub

Deep Decision Making and Reinforcement Learning: Final Project Submission

Team Members

Jayesh Chaudhari
jsc9903@nyu.edu		 Satyam Kumar
sk12075@nyu.edu		 Varad Vijay Suryavanshi
vs3273@nyu.edu		 Rivujit Das
rd3681@nyu.edu

Title

Online Exploratory World Model

Contents

1. Introduction

Reinforcement learning (RL) has demonstrated significant advancements across a diverse range of domains, from strategic board games to sophisticated robotic control tasks. Nevertheless, purely model-free RL approaches typically demand extensive interaction data, face considerable challenges in dealing with environments with sparse or long-horizon rewards, and necessitate substantial hyperparameter tuning and retraining for each new task even within the same domain.

To mitigate these limitations, world models have emerged as an effective paradigm. A world model learns an internal representation of environmental dynamics, enabling agents to anticipate the future states resulting from a sequence of actions. Such models provide the capability for the agent to internally simulate or “imagine” future scenarios, significantly reducing reliance on inefficient trial-and-error interactions with the real environment.

Historically, many world models have operated directly within pixel-space, reconstructing raw images to predict future observations. However, this approach incurs substantial computational costs due to intensive image reconstruction requirements and frequently relies on complex diffusion-based models. Consequently, recent research has increasingly favored latent-space prediction, wherein models operate on compressed, low-dimensional representations of environmental states.

Recent innovations like DINO-WM utilize pretrained visual embeddings such as those derived from DINOv2 to construct task-agnostic latent dynamics models. These pretrained embeddings eliminate the computational overhead associated with pixel reconstruction entirely, enabling more efficient prediction and planning. Moreover, models such as DINO-WM exhibit significant generalization capabilities across varying task configurations and environments, even in the absence of explicit reward supervision. This advancement underscores the potential of latent-space world models to address fundamental RL challenges, thereby advancing efficiency, generalization, and adaptability in reinforcement learning.

2. Proposed Methodology

2.1 Method 1 : Online Exploratory World Model

We try to work on these two recent methods and try to solve the issues in these methods. DINO-WM assumes having access to offline datasets with sufficient state-action coverage, which can be challenging to obtain for highly complex environments, and it also not the approach that humans would generally take while performing a task if you are play a game you would just know the basic rules or maybe not even that and start playing by taking random actions and making your understanding of the games dynamics better overtime. So we try to collect data using various exploration strategies. This data is mixed with optimal paths (reward maximizing actions) and suboptimal paths (for exploration) and trained on WM. These explorations strategies will involve maximising reward and reward would be of different types and would vary across environments.

Our data collection would look like following:

2.1.1 Actions Scorer

The reward strategies can be broadly categorized into extrinsic, intrinsic, hybrid, and hierarchical rewards. In our case the intrinsic reward strategies seem to be relevant so we try to work on them. We implement exploration/curiosity based reward strategies. Examples of these strategies in the pushT environment can be increasing the number of collisions between the pusher robotic arm and the T block, increasing pixel to pixel change in the environment per step.

Figure 1: Blue: optimal path (according to some action_scorer); red: suboptimal paths.

Figure 2: Action scorer.

In addition to global exploration strategies, we adopt a local exploration approach to effectively train and refine our world model. Specifically, starting from an identified optimal trajectory (represented by the dark blue line), we systematically explore additional nearby states within a defined local window. This local exploration involves investigating multiple alternative paths branching off from the current optimal trajectory. Within this local exploration window, we calculate and evaluate rewards for all potential state-action pairs explored. Importantly, if any of these alternative paths within the local window yield a higher reward compared to the previously identified optimal path, the superior alternative path (represented by the light blue line) replaces the current optimal trajectory, becoming the new focus for exploration. This dynamic updating ensures continuous refinement and adaptation of the optimal path based on the most rewarding outcomes discovered through local exploration. These locally explored suboptimal paths also provide diverse and valuable training data, enriching the world model’s understanding by covering a broader range of environmental dynamics. This comprehensive exploration methodology enhances the predictive capability of our world model, significantly increasing policy robustness and adaptability to diverse and unforeseen environmental conditions.

Figure 3: Tree based local search.

2.2 Method 2 : Ideal WM

The current approach in DINO-WM involves training the world model followed by planning, our proposed methodology integrates these stages into an iterative cycle. Initially, we perform an initial phase of world model training using exploration-derived data. Once the world model has acquired foundational dynamics knowledge, we proceed to a planning stage where optimized actions are computed. These optimized actions, derived from planning, are then incorporated back into further training of the world model, enriching its predictive capabilities and aligning its understanding closely with optimal decision-making patterns.

This iterative cycle consisting of alternating training and planning phases is repeated multiple times. Each iteration progressively refines the world model by continually incorporating the latest optimal actions identified during planning. This continuous feedback loop between planning and training ensures that the world model dynamically improves, effectively integrating strategic insights from planning into its predictive structure.

Figure 4: Existing Architecture

Figure 5: Proposed Architecture

2.3 Inducing Exploration in Dreamer-V3

To test our hypothesis on Dreamer-V3, we modify the agent’s policy method such that during environment interaction, the batch of actions is composed of both policy-driven and randomly sampled actions. The implementation details are as follows:

This ensures that in each training step, 50% of actions reflect learned behavior, while 50% inject purely exploratory behavior.

3. Results

Results PushT (Method 1 DINO WM)

The DINO WM results on the PushT environment reveal several limitations. Our use of greedy techniques to select optimal and exploratory actions fails to produce diverse or effective trajectories, which hampers the model’s ability to learn the environment’s dynamics comprehensively. While the greedy approach does encourage frequent interactions between the Pusher and the T block leading to effective learning of those specific dynamics, it neglects other critical patterns. For example, interactions where the Pusher approaches the T block, pushes it, then retreats and re-engages from a different angle are poorly represented and thus not well learned. To effectively replace the offline world model training proposed in the original DINO WM paper, we require a more robust action selector module capable of generating diverse trajectories that capture a broader range of interaction dynamics.

Results Atari (Method 1 DINO WM)

In the Atari Breakout environment, rewards are extremely sparse, making effective learning challenging. We attempted to capture meaningful interactions by inducing changes across different regions of the environment. However, both greedy and random exploration strategies perform poorly in this setting. The ball’s interactions are infrequent in the collected trajectories, as it often falls without meaningful engagement, preventing the world model from observing reward-generating behaviors. Additionally, tree-based reward selection fails to identify action sequences that align the paddle directly beneath the ball, which is crucial for success. Furthermore, other approaches that perform well on Atari games typically rely on training over a large number of epochs, a scale of computation we were unable to replicate.

Results PushT (Method 2 DINO WM)

PushT environment demonstrate incremental improvement over Method 1, with slightly more effective action behaviors emerging during planning. We can observe that interative planning and improving world model on the planned actions helps in improving the world model by covering tragectories that were not covered. Still in continues to suffer from from the same problem of non-ideal action selection as method 1. The absence of expert demonstrations continues to hinder the model’s ability to learn accurate physical interactions, emphasizing the need for better-informed action selection strategies and more diverse training data to improve model performance in complex physical environments.

Results (Inducing exploration in Dreamer-V3)

The primary reason our modified DreamerV3 model did not achieve the desired performance is the significantly reduced training duration. While the original DreamerV3 model was trained for 10^10 steps, our model was trained for only 10^5 steps, limiting its opportunity to thoroughly learn optimal policies. Additionally, introducing random trajectories as seed states for the imagination process inadvertently slowed policy convergence, as the model frequently imagined suboptimal or irrelevant scenarios. To address this, we propose masking these random and sub-optimal action instances during the imagination phase, ensuring the policy training focuses exclusively on trajectories derived from its learned distribution, potentially accelerating convergence and improving performance.


4. Conclusion

This project was our attempt to push the boundaries of exploration in world model-based reinforcement learning. Our methods introduced structured randomness through curiosity-driven rewards and suboptimal action paths. We tried to bring agents closer to how humans naturally learn such as by exploring, making mistakes, and gradually improving.

We worked with both DINO-WM and DreamerV3, testing new ways to collect data and influence learning. While the results were mixed, they gave us meaningful insights. In DINO-WM our exploration strategies failed to get good trajectories and thus it did not give very good results. We need a much better exploration strategy to replace the offline world model training process with an online world model. With DreamerV3, our modifications yielded some decent results on Atari environments. However, a direct comparison with the original DreamerV3 requires significantly more training epochs, making conclusive evaluation computationally expensive.

In the end, even though we didn’t hit perfect scores, this journey helped us better understand the balance between exploration and exploitation and how crucial smart exploration is for building better, more adaptable world models. There’s still a lot to improve, but we believe this work sets the stage for deeper investigations into learning more from less, especially in complex, open-ended environments.

5. Future Directions