Dear reviewer,

Thank you for your feedback and for noticing the generality of our method and the astute observation on the structure of the initial submission. We have updated our paper to ammend those concerns. Here we address the main questions in more detail:

Weaknesses

Figure 1 caption is added and the method sections has the pseducode to our Algorithm and color-coded to how it can be used to implement various UED meethods.

Our method section 3 has been updated with a pseudocode of The Policy Optimization with World model Ensemble Rollouts (POWER) algorithm. We use UED methods to sample the world models. We do not generate new ones in training time. The ubiquity of fully offline world model ensemle rollouts is that, by only changing the boolean flags, it can accommodate between different sampling methods both on the trajectory level and on the agent step level.

Empirical Evidence we have added additional D4RL results and Ablations for the size of the ensemble in section 5.5 and in A6.

Questions

1. Section 5.5 and in A6
2. Parallel computation runtime using our implementation in Appendix A3.

Wall-clock time in minutes on a single NVIDIA A40

4. While these benchmarks are outside the scope of our investigation and proposed method, we have no reason to believe in any inference time bottleneck given the results in the previous answer and our fully-jittable pipeline that is trained end-to-end in a GPU [1], [2]

[1] https://github.com/jax-ml/jax

Dear reviewer,

Thank you for the detailed feedback and the careful examination of our work. We appreciate the emphasis on our work's novelty, lack of imposed constraints on the policy training and the relevance of our work to the current state of RL research. In addition to the clarifications in the global response, here we target your specific conerns.

Weaknesses

1. We have added the explanation of every symbol in the equation. We introduce the approximation of regret to make our method more general as the optimal performance in a given world model is not always known. We have added the pseudocode for our Algorithm. We have improved the caption in Figure 1 to make sure both sides are clear.

2. We have shortened the preliminaries and focused on making the method easier to parse. We added a line in the CMDP prelims clarifying the connection.

3. We agree that the claim of D4RL being adversarial is not supported by the distribution of the dataset and recent work cited in this paper.

4. The Figure is an illustration of the sampling process, we try to change the sampling frequency to capture as much of the learning phase as possible and discard after a handful of converged checkpoints. We have made sure to have that reflected in our writing.

5. We use MLP as those are used in offline-RL methods and lend themselves more naturally to parametric levels in the UED literature.

6. Our method is not PLR. We bring together two aspects of Reinforcement Learning that can remove their individual limitations, namely the hand-crafted domains UED is mainly used for and the explotability of simple world models. Our method is now more clearly described in the method Section 3 and Algorithm 1 that allows for the use and testing of different curricula curation methods.

7. We would like to clarify that we have always used the ICLR 2025 format. We have removed any additional setting that would change the default color of the links.

Questions

1. The maximum size is 1 million and we downsample from that at the transitions level. Full table in appendix A2.

2. Yes, provided in Section 3 of the revised PDF.

3. Our implementation is fully based on CORL [1] and its JAX variant [2], matching their reported performance benchmarks. For fair comparison, we downsampled datasets at the episode level rather than transition level when evaluating CQL and SAC_n. Our results fall within the 95% confidence interval of the provided implementations. We developed SAC_N and CQL on a shared framework for consistent evaluation, and will release these baselines alongside our method. Our implementations achieve match the reported scores across all MuJoCo tasks.

4. Apologies for the oversight. All has been corrected.

5. You assessments of the general tendencies in each subfield is correct. We are addressing two notable limitations of offline RL which is that of truncations and biased datasets and simultaneously point out the benefit of online RL breakthroughs. We use present UED and our world model ensemble paradigm to not only oversome the common model-based offline-RL "pathologies" but also show how online RL can work in general abstractions of the real world without hand-crafted contraints.

[1] https://github.com/tinkoff-ai/CORL

[2] https://github.com/nissymori/JAX-CORL

We appreciate the reviewer's positive feedback on our revisions and the score indrease. However, we respectfully disagree with the suggestion to compare with pixel-based settings, as these scenarios fundamentally differ from the challenges our work addresses. Our research specifically focuses on solving efficiency problems in dynamic models, and our experimental design directly demonstrates the effectiveness of our proposed solution.

We have provided evidence through:

1. Clear experimental results demonstrating our method's capabilities
2. Additional D4RL benchmark results
3. A detailed implementation of our computationally efficient pipeline

Given that the reviewer acknowledges the improved clarity, we would appreciate further clarification on why the absence of pixel-based comparisons and additional SOTA benchmarks (which lie outside our paper's scope) would lead to such a low score. If you feel strongly that these comparisons would significantly strengthen the paper and are the only ground for rejection, we are willing to include these additional experiments for the final version.

Dear reviewer,

Thank you for your feedback and for higlighting our combination of UED and world models, lack of reliance on online tuning and improvement over the use of a single world model.

Comparisons

We use a simple MLP as done in the model-based methods like MOPO. Those methods have very specifically tuned dynamic models for each environment and dataset. MOPO and other model based methods either truncate in 3 to 5 environment steps or diffuse the entire trajectory like PGD [1]. Model-based offline RL methods train an ensemble and only use it for inference during the short model rollouts.

Our method does not require turning the number of ensemble member and other hyperparameters related to using ensembles. We use every member of the ensemble as a separate level in the Unsupervised Environment Design paradigm. Therefore, we use the same architecture for every model and do not need any hyperparameter tuning to match the MOPO results in tasks like Hopper. The costly and specific tuninf per-task is reported in notable implementations too [3] [4] [5].

Minor concerns

Figure issues

The returns of the behavior policy trained to convergence is shows with a red line for all the Gymnax and Brax environment. We have added clearer online PPO baselines in the D4RL results.

Use of terms

We have changed the clarify the following: world model is referred to our training, real world refers to the real environment Gymnax or Gym environment where the data is collected and the policy is evaluated. This misuse stems from the fact that the MuJoCo environments are physical engines built to stimulate rigid robotic interactions observed in the real world. This is a very important distinction we are thankful you addressed.

Appendix

The then A3 not A5 Hyperparameters section content was after the large table with the test losses for the world models due to ICLR template rules. Apologies for the oversight, we have rearranged the appendix to match the location of each heading with its subsequent content.

[1] Jackson et al. Policy Guided Diffusion

[2] Sims et al. The Edge-of-Reach Problem in Offline Model-Based Reinforcement Learning

[3] https://github.com/tinkoff-ai/CORL

[4] https://github.com/nissymori/JAX-CORL

[5] https://github.com/yihaosun1124/OfflineRL-Kit

Thank you for discussing another model based RL method like MOPO and for addressing some of the minor concerns. I think it is important to see performance of the proposed method on image based environments like Atari as well as compare results with the recent improvements in world models such as in IRIS [1] (which also uses the same fixed set of hyperparameters across atari environments). I am happy to raise my score, if the authors can provide results on image based environments and possible comparisons to established world model papers.

[1] V Micheli, et al, Transformers are Sample-Efficient World Models, ICLR 23

Dear Reviewer,

Thank you for qualifying our work as a fresh perspective on RL and our demonstrated robustness under a data-efficient regime compared to conventional offline methods. Here we address specific points in addition to improving the structure and readability.

The world models term

Your observation regarding the terminology of our method is correct in that we do not require a vision encoder or a recurrrent component to make predictions on such encoding. We use the term world model because we are using the trained model as an abstraction of the real environment by directly predicting the next state and reward given the action at given state. Traditional offline model-based methods like MOPO [1] use dynamic models to predict changes in the observation instead of the next state changes and incorporate uncertainty penalties. Our algorithm deliberately assigns a single model for each single-step prediction, in contrast to conventional model-based implementations [2] that use ensemble inference methods -- like greedy or a linear function of the top k elites -- for each observation. This in turn introduces more hyperparameters that our world model method does not require. With this distinction made, we are kindly open to using a different term like predictive, dynamic or parametric environments if you deem it necessary for accuracy.

Why UED

As we further explain in section 3.2 and Appendix A1, UED has several desirable properties in our case. Its empirical success and theoretical guarantees can help the agents generalize to test distributions that were not seen during training but are in the support of the design space i.e. our $\Theta$ space of potential parametrizations containing the real MDP $\theta^*$ that is highly unlikely to be sampled during training as no world model can perfectly fit the real MDP i.e. the real environment. Also, by using UED to select different world models during the training process, we mitigate catastrophic exploitaion of a single model as shows in Section 5.1 which would undermine the transfer of the agent anywhere outside the exploited world model (level).

D4RL dataset

D4RL results competitve with online PPO in sec 5.4, ensemble size ablation in sec 5.5 and Appendix A.6.

Computation Overhead

We note that as described in the Algorithm, the World Models are trained using the same data and vmapped and JIT-ed in a single Nvidia A40. 100 is chosen to show the schalability of the method and to match the large design space in the UED literature [3] [4]. While our ablations show that the method is robust to smaller number of world models, our algorithm allows for a combinatorially large number of levels to be generated by training on the same data. We hope for our method to be used beyond the current horizon of salient RL problems.

We note our PureJaxRL-based RL training [5] has 100x speedups to standard PyTorch baselines. Appendix A2 and A3 show our wall-clock efficiency using vectorized mapping vmap

[1] T Yu, et al, MOPO: Model-based Offline Policy Optimization, NeurIPS 20

[2] https://github.com/yihaosun1124/OfflineRL-Kit/blob/main/offlinerlkit/dynamics/ensemble_dynamics.py

[3] Dennis et al, Emergent Complexity and Zero-shot Transfer via Unsupervised Environment Design

[4] Jiang et al, Prioritized Level Replay

[5] https://github.com/luchris429/purejaxrl

Dear Reviewer,

We would like to thank you for your appreciation of our work and for noting how it introduces a novel approach to online RL and our attempt to mitigate common pitfalls of offline methods. In addition to fixing the typos, D4RL experiments and ablations that are further discussed in the general comments, we address your questions:

1. Grid Search for baselines

Our implementation builds upon the CORL codebase by Tarasov [1] and its JAX port [2]. We carefully matched the implementation details and hyperparameters to reproduce the highest reported performance from [1], [2] and their respective baseline methods. After validating our reproduction, we conducted a grid search over the hyperparameters for each of ratio of our dataset. Even though our method only requires single-step transitions, we maintained a fair comparison with CQL and SACn by downsampling episodes uniformly rather than individual transitions, as both CQL and SACn were designed to operate on complete trajectories.

2. D4RL Evaluations

We agree that more benchmark evaluations are necessary to validate our Algorithm and method outlined in Section 3. Claiming that current offline RL methods exploit safety biases in common benchmarks by introducing a dataset where they do not perform as well in comparison does not absolve us of the responsibilty of testing our method beyond our own curated dataset. We have added D4RL results in Section 5.4 using the full-replay datasets as they match the $10^6$ transitions magnitude used for the other experiments. We achieve results that are competitive with the online PPO results in the respective environments [3]. A table can be found in the global response.

3. Open Source

We will open-source the code used for all the experiments with the WandB integration with an abstract class template for any parametric environment for full length rollouts.

[1] https://github.com/tinkoff-ai/CORL

[2] https://github.com/nissymori/JAX-CORL

[3] https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/

To the best of our knowledge this is simply not right, all the papers we have looked at use truncated rollouts to prevent compounding errors inside the world model. Specifically, MoReL applies a large negative reward and truncates rollouts when there is a discrepancy between the different dynamics models in the ensemble. If you find a methods which do indeed use full-length rollouts we would love to see this. In the absence of such evidence we find this claim rather unfortunate as a basis for the rejection of the paper.

Furthermore, this should have been brought up earlier in the review process to allow us responding to it.

The Authors.