Wasserstein Robust Reinforcement Learning

Abstract: Reinforcement learning algorithms, though successful, tend to over-fit to training environments, thereby hampering their application to the real-world. This paper proposes WR2L -- a robust reinforcement learning algorithm with significant robust performance on low and high-dimensional control tasks. Our method formalises robust reinforcement learning as a novel min-max game with a Wasserstein constraint for a correct and convergent solver. Apart from the formulation, we also propose an efficient and scalable solver following a novel zero-order optimisation method that we believe can be useful to numerical optimisation in general. 
      We empirically demonstrate significant gains compared to standard and robust state-of-the-art algorithms on high-dimensional MuJuCo environments

Reviews and author responses:

The paper studies the problem of training policies which maximize returns under worst case environmental perturbations. They propose to solve this max min problem using iterative optimization with the max step performed using standard RL. The min step relies on sample-based estimates of gradients and hessians with respect to environment parameters. The problem of robust control is an important one, and the paper presents interesting results, however it lacks a proper comparison to baselines.

1. I think figure 1 should include two more columns. The first showing the performance of a policy trained to maximize returns averaged over the training environmental perturbations. This doesn't require the min step of the algorithm and is a natural baseline for the setup you consider. The second showing the same but trained with an LSTM policy. This way the policy can learn to do system identification. It might not be robust (i.e. not generalize to unseen perturbations) but this should be verified. If it does generalize, then there's no reason to complicate things with the min step.
2. Maybe I'm missing something but neither RARL nor PR-MDP baselines are shown in Fig. 1 and 2. despite you claiming in intro to section 5 that you test against those. These baselines are included in Figure 3 in the Appendix but I find those numbers odd. E.g. RARL doesn't seem to match PPO even on the reference environment in Fig 3c. This suggest that it relies on a different algorithm and it doesn't make sense to compare it to WR2L which relies on PPO.
3. In Related work you says "we were unable to find the code for this paper, and did not attempt to implement it ourselves". Even if this code was available, comparing to a different codebase is not a meaningful comparison (see e.g. [1]). There is no replacement for implementing relevant baselines.
4. You estimate the Hessian using data from a random policy. I find this odd because random policy never visits the interesting states. Have you tried recalculating the Hessian in every min step?
5. Could you clarify the use of Wasserstein distance? If I only use samples from the two distributions to estimate it then this is likely limited to low-dimensional settings, right?

[1] Henderson et al., (2017). Deep Reinforcement Learning that Matters.

------------------------------------------------------------------------------------------------------------------------------------------------------
Thanks for the clarifications and for including the "average" policy baseline. I still think the paper could use more work and hence keeping my score as is. I'd try to make the LSTM sys-id baseline to work--those have been used with domain randomizations previously, as well as "meta-learners" in the meta RL literature.

1) Keeping the original formulation would mean an (uncountably) infinite number of constraints, one for each state-action pair. Each of these constraints would in turn be a highly complex, non-convex function which maps simulator parameters to Wasserstein distance. 

How would the reviewer suggest this problem be dealt with without any relaxations or approximations? 

In any case, as we see from the experiments, the performance doesn't fall off much when compared against  PPO trained on the reference dynamic (which is the best we could expect). 


2) As we stated, our paper is focussed on addressing the problem of simulators not being accurate representation of real-world dynamics. This is a major stumbling block for RL. Reward mispecification is an important problem, but it's a different one, related more to the dificulty in encoding the desired behaviours into one-step reward functions.  We do not believe that the solution to the reward mispecification problem would necessarily look like the solution to the dynamics mispecification problem. They are a different class of mispecification.

We have already cited and compared against other published papers in that are focussed solely on the problem of an imperfect simulator dynamics in Deep RL. To say that our paper is incomplete is to say theirs is too. 

Response to Author Response: 

Thanks for authors’ detailed explanations. It is an interesting topic to solve the deep robust reinforcement learning problem. However, I still have the following concerns of the paper. 

1) The relaxation of the Wasserstein constraint in eq.(7) indicates an approximate solution to the original robust reinforcement learning problem, therefore, is conservative. 

2) Considering existing works mostly handling both reward and transition dynamics, it is incomplete that the paper does not provide a solution for the reward dynamics.


We thank the reviewer for their time and valuable suggestions. 

1. We have run new experiments and added a fourth column to Figure 1 and Figure 2 to demonstrate effectiveness of the reviewer’s suggestion for random environment dynamics. The results are strong in the two-dimensional case (Fig 1), but clearly inferior to our method in high dimensional settings (Fig 2). 

On LSTMs, we also took the reviewer’s advice and attempted some experiments, but we saw the performance was poor. Others  have tried as demonstrated here https://github.com/openai/baselines/pull/859 where it can be seen that the scores achieved are struggling to reach scores of 2500, which our algorithms are easily beating.

2. The results are, as the reviewer mentions, in Figure 3 in the Appendix.  Regarding why RARL doesn’t do as well as PPO in the reference environment, this is entirely to be expected, since RARL aims to be robust. It would be quite strange for an algorithm that aims to be robust to environment variations to do as well as PPO trained to be optimal on one particular environment. Indeed, we expect PPO to be better on the reference environment than our own algorithm. 

3. In fact, following the reviewer’s advice, we implemented the PPO trained with mean dynamics and added the results to our paper. This setting is similar to Rajeswaran et al when the \epsilon variable in their paper is set to 1. As we see, it does well in low dimensions, but poorly in high dimensions. From the paper of Rajeswaran, we see that other values of \epsilon (the \epsilon in their paper) the results can be somewhat better but not much. So whilst we do not claim to have implemented and tested their algorithm in full, the new experiments give us some idea of its strengths and weakness. 

We note that this is entirely what we expected: The Algorithm if Rajeswaran et al implements the minimisation search by sampling dynamics at random and selecting some worst-performing fraction of them to update policy. This strategy will be fine for low dimensions but in optimisation in high dimensions through random sampling is ineffective because of the well-known curse of dimensionality. 


4. Yes, we tried recomputing the Hessian, but the problem is that one ends up with an optimisation problem with a moving constraint set, which is difficult to deal with and did not work much better than our current method.

It is also not false to say a random policy does not visit interesting states, a random policy will visit all states that it is possible to visit, which is why it is the standard method of exploration in e.g., Q learning. 


5. This is a good point. Our dynamics are treated as Gaussians which makes their computation easier in high dimensions because there is a closed-form expression for the Wasserstein distance between Gaussians based on mean and covariance. For general distributions, estimation of Wasserstein distance could be difficult in high dimensions; at the moment we do not know the best way to handle this problem. We state however, that we chose to test on Mujoco environments that are standard benchmarks – we did not seek out easy cases. 

We thank the reviewer for their time in looking at our paper, and we appreciate their feedback. 

1) We wish to be very clear on this point: Our approach is not model-based RL. We do not try to learn the transition model, so it is not a case of system identification. We explain this at the end of the first paragraph, but to reiterate, the set-up is like this: You have a simulator of some environment/system (e.g., robot) that you would like to learn to control. The simulator is parameterisable (which is the case with the extremely popular and standard Mujoco benchmarks). Our algorithm then gives a policy which is robust to variations in dynamics parameters after the training is complete. The long term goal is to have a policy which is deployable in environments for which the simulator is not a perfect representation. Whilst the algorithm takes gradient steps in dynamics space, at no point is there explicit model-learning. 


2) We have rewritten the explanation of the sampling, under Eq. (7), to be simpler and hopefully more clear.  The reviewer is correct in saying "..Is it sampling data from the P_\phi and \pi on one hand and P_0 and \pi_{Uniform} ", and further, the emperical distrubution so generated is used to estimate the Wassertein distance. This can be done using linear programming or solving min cost flow problem, but we leave that to a reference like Peyre and Cuturi, (2019)

The "bucketing" terminology has been removed. Hopefully the new paragraph is clearer, but essenitially, the Expectation is over (s,a) pairs that are sampled uniformly from collections of trajectories, that themselves are generated from  P_0 and \pi_{Uniform}. 

Note that, as the text following explains, these are used to determine H_0. 

3) The main approximation is in number of emperical samples for approximating the Wasserstein distance. One possible aproach, which worked well in our  experiments, was to assume that the the dynamics are given by deterministic functions + Gaussian noise with diagonal convariance matrices. This makes estimation easier in high dimensions since sampling in each dimension is independent of others, and the total samples needed is a constant factor of the number of dimensions. Gaussian distributions also have closed form expressions for Wasserstein distance, given in terms of mean and covariance. 

We thank the reviewer for this point and have updated our paper to include this information. 

4) All the methods we compare – our own and others – have access to the same information and functionalities. They all have access to a simulator which they use to learn a policy. They all have access to parameters of the simulator (that being, the dynamics parameters), but how they use those parameters (if at all) is different. 

When we train using a particular algorithm, we give it the simulator with some default settings, and allow it to train however it is designed to. The output is a policy. So each algorithm will return a policy. We then  take these policies and deploy them on the simulator again, but we run multiple experiments – each experiment with a different setting of the dynamics parameters. We make sure that the algorithms are tested against parametersthey did not see during training, to test how robust they are. 

As can be seen from figure 1, where we compare our algorithm to PPO, the performance of our algorithm is fairly even across different dynamics, whereas PPO does well in the region close to where it was trained, but performance falls off for dynamics which are more different. A similar conclusion can be drawn from the figures in the appendix, where our algorithm does better than others. 


5) In fact, this is the whole point of our paper, that the available simulator is not a perfect representation of the system, and we want an algorithm that can cope with this imperfection. Remember, we do NOT have a model of the system (we can try to learn a model, but this is a different problem). All we have to train our policies is an imperfect simulator. 

The caveat is that the simulator needs to have parameterisable dynamics, and those aspects of the dynamics which are not parameterisable our algorithm may not be robust to, since it can’t learn to be robust to them. 


6) In fact, learning models of dynamics is a future reserch direction for us. The problem is that model learning is a highly non-trivial problem when dealing with high-dimensional state spaces. State of the art methods apply Model Predictive Control (MPC) or closely related. 

Moreover, varying model parameters can lead to bizarre situations with implausible physics, such as non-Newtonian movements. So even if we had a perfect model, if policy is to be robust, we will need to find a way to vary dynamics in reasonable way, otherwise we may as well just use PPO which is very effective in high dimensions. 




We thank the reviewer for their time and consideration, and are grateful for their feedback. 

In response to your points,

1) The reviewer is correct that the work of Yang (2017) already formulated robust policy learning as a Wasserstein constrained optimisation problem; we tried hard not to give the impression that we were the first to do this. Indeed, in the second paragraph, we cited Yang (2017) and state:

“…Whilst we are not the first to introduce the Wasserstein distance into the context of MDPs (see, e.g., Yang (2017) or Lecarpentier & Rachelson (2019)), we believe our formulation is amongst the first suitable for application to the demanding application-space we desire, that being, high-dimensional, continuous state and action spaces…”

Furthermore, in the Related Work section, we twice emphasised the work of Lecarpentier and Rachelson (2019), which is even closer to our setting.  Our novelty lies primarily in a method which effectively learns robust policies in high-dimensional state and/or action spaces. 


We would argue against the claim that our work is incremental. What we have presented in our paper is:
- a Deep RL method that… 
-…demonstrates effective learning of robust policies… 
-…in infinite (“continuous”) state and action spaces, in particular,…
-…those of high dimension. Furthermore, 
-…it is a method which appears to do better than other methods, at the time of writing. 

Yang’s work does not attempt to do this, as it is based on MDPs with finite state and action spaces. Methods that are effective on finite state and action spaces often do not carry through to large state and action spaces, particularly those of high dimension. 

What we meant with “… a novel min-max game with a Wasserstein constraint…” is really that we formulated an optimisation problem (in particular, eq. (9)) which is amenable to Mujoco-type problems that we test. However, we take on board the reviewer’s assertion that this claim of novelty is over-claimed, and so we will change 

“…Our method formalises robust reinforcement learning as a novel min-max game with a Wasserstein constraint for a correct and convergent solver. Apart from the formulation, we also propose an efficient and scalable solver following a novel zero-order optimisation method that we believe can be useful to numerical optimisation in general….”

to

“…Our method formalises robust reinforcement learning as a min-max game with a  constraint derived from Wasserstein distance,  and we propose an efficient and scalable technique algorithm to solve it. Our algorithm applies zero-order optimisation method that we believe can be useful to numerical optimisation in general….”

We have also changed the Conclusion . 



2) Ambiguity in the reward function is an interesting avenue of possible investigation when considering other aspects of RL robustness, and the Wasserstein constraint is a plausible constraint in this setting too. Indeed, the formulation of Yang (2017) includes reward ambiguity . However, for the current paper we wanted to keep a tight focus on the dynamics aspect of ambiguity, and felt that another type of ambiguity would be suitable for an extended version of the paper, e.g., for a journal. 


3)Quite possibly, yes. Since the condition now only needs to hold on average, it is possible that it may be violated for some state-action pairs. This translates to the policy being judged against a worse set of dynamics than in the preceding constraint formulation. We will note this observation in the revised version of the paper. We thank the reviewer for this point. 

Regarding the effect of epsilon, we do not yet have a definitive answer,  but some indication is given in Figure 1 where we see that increasing \epsilon leads to a more uniform performance of the learned policy across variations of dynamics parameters (friction, density, etc), but that the colours are lighter – i.e., the highest return is reduced. This is to be expected, since being robust to a wider range of dynamics intuitively suggests that the policy would not be as good as one trained on a more focused set of dynamics. 


4) Thanks for spotting this omission. Darker is better, with the numbers in the graded scale being the range of scores that the tasks can generate. We will add this explanation to the paper. 


The paper proposes a novel solution to a *restricted* reinforcement learning problem. I am not quite sure if I understand the restriction from the text or thus what relevance this has in applications. The method factorizes the transition dynamics (where they rely on a parametrized solver) and a policy modeled by a NN and trained with policy gradient with monte-carlo returns. The ultimate goal is robust control so they care about the maximin objective of highest return along the worst trajectory. The robustness is implemented as a constraint in the dynamics optimization: a wasserstein distance between the distribution of samples obtained applying the learnt policy and learnt dynamics parameters and those obtained by random policy and some reference dynamics model. They solve the maximin by alternating minimization on the dynamics and policy. The contribution is in the intricacies of computing the wasserstein efficiently in particular they obtain a perturbation based estimate of the gradient of the dynamics parameters. The experiments seem sensible but I  have some suggestions for improvement.

I have published in RL but my familiarity with Wasserstein and robustness is limited. I have not checked the derivations in the appendix in detail. If there is a high score I'll have a closer look. 

Major points:
1) The restriction imposed is that only part of the transition model is really learnt. So it is more like a system identification setup ?! Why is that necessary ? 
2) The explanation in on page 4 about how the wasserstein distance is computed is a bit unclear. Is it sampling data from the P_\phi and \pi on one hand and P_0 and \pi_{Uniform} ? Then there is a bucketing step (what does it mean ? how is it computed ?). Is it then a discrete wasserstein between the two empirical distributions ?! What is Monte-Carlo in this case ?
3) All the approximations done in the method should be more clearly explained because they are relevant both to people who might apply the method in practice as well as other researchers that might want to build on/ improve on it.
4) The experiments seem interesting but I am not sure how to interpret the results. If a method has access to a simulator and a representation of the relevant dynamics parameter space it seems clear that it will be more robust. This is impressive if one has a simulator and can infer some state components e.g. in pendulum starting with an empty simulator and adding a pendulum. But that seems as hard or maybe even harder than trying to learn an implicit transition model just from observations. 
5) Imagine the method only gets observations and an imperfect simulator, which will be the case in an application, what are the caveats we should have in mind ?
6) Can you have a learnt model baseline and compare your learnt and simulation ? Maybe it would have to be a latent variable model so you can do the same inference for internal parameters and still do an adaptation in the same way as now. This seems easy since it is a supervised problem.

This paper focuses on the problem of robust reinforcement learning and proposed a generic framework that is able to deal with both discrete and continuous state and action spaces. They formalized the robust reinforcement learning problem as a min-max game with Wasserstein constraints and put forward an efficient and scalable solver that is helpful to numerical optimization. Overall, this paper is of well-written with a clear illustration of the methodology and comprehensive experimental results. I still have a few concerns below, that prevent me from giving a direct acceptance: 
1.	However, the proposed method seems to be somewhat incremental: In Yang (2017), it is clearly stated that it is viable to treat the robust reinforcement learning objective as a min-max game (in eq.(1)) with support belonging to a Wasserstein ball (in eq.(2)). The statement in the submission “… a novel min-max game with a Wasserstein constraint …” seems overclaimed, namely, eq.(8) in the submission, is a combination of eq.(1) and (2) in Yang (2017), with the constraint in eq.(8) not exactly equal but just a relaxation of eq.(2) in Yang (2017).
2.	Furthermore, since a few papers have proposed methods to deal with both transition and reward dynamics, it would be nice and complete to also address (or hint on) the reward ambiguity problem.
3.	Does the relaxation of the Wasserstein constraint in eq.(7) make the learned policy conservative? Some illustrations of this effect of \epsilon would be interesting and complete.
4.	For the experiments in section 5.1, it would be better to give an illustration of what the color means in Figure 1. 
Insoon Yang. A convex optimization approach to distributionally robust markov decision processes with Wasserstein distances. IEEE control systems letters, 1(1):164-169, 2017.


Meta-review: While the reviewers agreed that the problem of learning robust policies is an important one, there were a number of major concerns raised about the paper, and as a result I would recommend that the paper not be accepted at this time. The important points are: (1) limited novelty in light of prior work in this area (see R2 and R3); (2) a number of missing comparisons (see R2). There is also a bit of confusion in the reviews, which I think stems from a somewhat unclear statement in the paper of the problem formulation. While there is nothing wrong with assuming access to a parameterized simulator and studying robustness under parametric variation, this is of course a much stronger assumption than some prior work on robust reinforcement learning. Clarity on this point is crucial, and there are a large number of prior methods that can likely do well in this setting (e.g., based on system ID, etc.).

Decision: reject