Operator World Models for Reinforcement Learning

Summary

This paper extends policy mirror descent to the RL setting. They solve the problem of requiring a exact value function by formulating an approximate value function based on operators over the transition and reward functions, which yields a closed form solution. Then they use this approximate value function for policy improvement. Lastly, they provide theoretical analysis on convergence guarantees and error bounds while explicitly laying out assumptions. They also perform experiments on toy RL problems and show that their algorithm outperforms typical RL baselines on these low-dimensional problems.

Strengths

1) This paper's presentation is great. How this work relates to prior work is clearly discussed. Key concepts from this paper are well explained and well motivated. 2) Theoretical analysis is sound and shows promising results. 3) Comparison to RL baselines, even on toy problems, is much appreciated. Results are encouraging. 4) Discussions of limitations and assumptions are detailed. 5) Key results of this paper, IE value functions with closed form solutions via operators, are novel and significant.

Weaknesses

1) I would like to see experiments on more complicated MDPs, with large state-action spaces, to test the scaling properties. I recognize that this paper already makes large contributions, so this is not necessary for now, but it is a great future direction. 2) I would like to see extensions to continuous action spaces. Again not necessary for now, but it seems like it should be possible given the formulation with additional approximations. Minor: Typo on line 335

Questions

How well can this algorithm scale to higher-dimensional settings?

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Authors have adequately addressed limitation and assumptions. I see no potential negative societal impacts as a result of this work.

Summary

This paper proposes an approach (the first PMD approach adapted to RL setting) that can learn a world model using conditional mean embeddings (CME). The operatorial formulation of RL is used to express the action-value function in closed form via matrix operations. The proposed algorithm is proved to converge to global optimum, and validated by simple/preliminary experiments.

Strengths

To the best of my knowledge, this work is the first one to apply PMD to RL setting with provable global convergence guarantee. The references and previous work are cited and well-discussed. The theoretical contribution of this work is sound and solid. The proposed method is validated to be more sample efficient than the current state-of-the-art RL algorithms like PPO, DQN, TRPO

Weaknesses

Line 88: "Borrowing from the convex optimization literature", please give the exact reference. Line 97: at the end, lack of a "space" From Line 91 to the end of the manuscript..., it seems like the hyper-link of references does not work. They lead me to incorrect place, and make it hard to follow the idea. The experimental validation is limited to simple/toy environments in OpenAI Gym. In the caption of Figure 1, I would change "dark lines" to "solid line". The authors should also mention the reward threshold is the "success" reward in the caption, not just in the main texts. It is probably better to add the lable of dashed line in the plot. Since the Figure 1 is in logscale, the authors should show the learning curves starting from 0 or 10^1, so that we can observe the full learning process.

Questions

In Figure 1, it is a bit surprising that the variance of the proposed method in (a) and (b) is 0 (or almost 0). I know that general RL algorithms will still have some fluctuation. Maybe the authors can give some intuition why the proposed method achieves such small variance? In Figure 1(c), why the proposed method stops at 10^4 step? I would expect to see the full learning curve, otherwise it is not convincing that the algorithm is convergent. Does POWR work in continuous action space? Even simple tasks in Gym including Pendulum, MountainCarContinuous, LunarLanderContinuous, etc.

Rating

7

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

The major limitation of this work is the empirical validation, where the authors only consider three toy/simple environments in OpenAI Gym. Also, the experiments are limited to the discrete action space. It will be better to see if the proposed algorithm works in continuous action space like robotics, etc.

Summary

(This review has been updated after revert of desk reject)

Strengths

(updated after revert of desk reject) It's a solid theoretical paper with nice writing and experimental results on continuous control benchmarks.

Weaknesses

(updated after revert of desk reject) The experimental results and more interpretation/justification of the results are limited. It would be stronger if the authors could further connect the experiments with the theoretical results.

Questions

N/A

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

yes

Summary

The paper presents a practical implementation of Policy Mirror Descent (PMD) for Reinforcement Learning (RL). PMD requires knowledge of the action value function of the current policy at each iteration. Existing methods that approximate the action-value function depend on the ability to restart the Markov Decision Process (MDP) from each state multiple times. Instead, this paper proposes learning estimators of the reward function and the transition model and combining them to approximate the action-value function. To facilitate this, an RL formulation using operators is introduced, enabling the use of conditional mean embeddings for approximation. The authors study the convergence of the proposed algorithm to the global optimum and compare it with other RL methods on classic Gym environments.

Strengths

- The paper is well-written, and the authors thoroughly introduce their method in relation to existing techniques.

Weaknesses

- Considering the connection between the proposed method and world models, the authors should discuss world models more thoroughly [1-3]. [1] Schmidhuber, J. Making the world differentiable: On using supervised learning recurrent neural networks for dynamic reinforcement learning and planning in non-stationary environments. (1990) [2] Schmidhuber, J. An on-line algorithm for dynamic reinforcement learning and planning in reactive environments. (1990) [3] Ha, D., and Schmidhuber, J. World Models (2018)

Questions

See above. The paper is rather technical, and several sections are beyond my expertise. I am looking forward to discussing it further with other reviewers to potentially increase the score for this submission.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are discussed in the paper.

Rebuttal

Thank you for providing the updated Figure 1. It indeed looks more reasonable. But is there a reason not to run 10^6 timesteps for MountainCar? I understand that rebuttal has limited time, but MountainCar is a very easy problem, which should not take that long. Will the authors plot 10^6 timesteps in the final version of the paper?

Official response

Thank you for your response. Please note that the claim "The notion of world models for RL was introduced by Ha and Schmidhuber in [R3]" is incorrect. World models in RL were already introduced in Reference R1 (1990) and in Dyna ([4], 1991). I trust that the authors will accurately discuss the background of world models. Having read the other reviews and the authors' responses, I am happy to increase the score for this submission. [4] Sutton, R., Dyna, an integrated architecture for learning, planning, and reacting. (1991)