Model-based reinforcement learning (MBRL) aims to learn a dynamic model to reduce the number of interactions with real-world environments. However, due to estimation error, rollouts in the learned model, especially those of long horizons, fail to match the ones in real-world environments. This mismatching has seriously impacted the sample complexity of MBRL. The phenomenon can be attributed to the fact that previous works employ supervised learning to learn the one-step transition models, which has inherent difficulty ensuring the matching of distributions from multi-step rollouts. Based on the claim, we propose to learn the synthesized model by matching the distributions of multi-step rollouts sampled from the synthesized model and the real ones via WGAN. We theoretically show that matching the two can minimize the difference of cumulative rewards between the real transition and the learned one. Our experiments also show that the proposed model imitation method outperforms the state-of-the-art in terms of sample complexity and average return.
This paper considers the model-based reinforcement learning (MBRL)
problem in which one of the main parts is to learn the transition
probability matrix of the underlying MDP (called Transition Learning -
TL). The motivation is that if the transition model can be learnt from
some real-world trajectories, the agent can improve by interacting with
simulated environment built from the learnt transition model; hence,
the overall number of real-world interactions is reduced.
Given a policy $\pi(a|s)$, and a transition $T(s’|s,a)$, the occupancy
measure $\rho(s,a)$, defined in Eq (2) in the paper, can be
interpreted as the discounted distribution of state-action pairs in
the rollouts. The main idea is to use Wasserstein-GAN (WGAN) to match
the distribution of $p(s,a,s’) = \rho(s,a) T(s’|s,a)$. In particular,
given a policy, the idea is to get some real-world trajectories and
feed into the WGAN framework in which the generator tries to generate
synthesized trajectories and the discriminator tries to discriminate
between them. These trajectories are then used in policy optimization
to obtain a new policy.
The authors presented two theoretical results: 1) Consistency for
WGAN: if WGAN is trained optimally, then the synthesized transition is
the same as the real transition; and 2) Error bound for WGAN: the
difference between cumulative reward under synthesized transition and
that under real transition is upper bounded by a linear function of
WGAN training error. On the experimental side, the authors used the
model-based benchmark library (MBBL, Wang et al. 2019) to compare the
proposed algorithm with several existing algorithms on four MuJoCo
tasks: Hopper, HalfCheetah, Ant, and Reacher.
Overall, I think this work is positive. The theoretical part is
technically sound (there might be an error in the proof of Proposition
1 but it is fixable). The experimental part is also sensible, although
it would be good to include the performance comparisons for other
tasks in MBBL. Details in the comments below.
Main   comments/suggestions:
- MBBL (Wang et al. 2019) has many other tasks (18 in total, the
  authors only include 4 in this work). It would be good if the
  authors can also include the performance of the proposed algorithm
  with respect to those tasks.
  
- The claim at the end of Proposition 1’s proof (in side the proof), in my assessment, is not
  established. It is clear from the proof that $p(s,a,s’) = p’(s,a,s’)$ is
  a sufficient condition for $\rho_{T}(s,a) = \rho_{T’}(s,a)$, but I’m
  not convinced that it is also a necessary condition. Note that
  $\rho_{T}(s,a)$ and $\rho_{T’}(s,a)$ are two unique solutions of two
  *different* Bellman equations. It would be good if the authors can
  provide the detailed proof for the necessary condition if it is
  true. Nevertheless, the statement of Proposition 1 only claims the
  sufficient condition so this is fixable.
  

Other minor comments:
- End of Section 1, page 2: the notations $T$, $T’$, $R(T)$, $R(T’)$ are used but not introduced until Section 3.
- Typo, page 2, last paragraph: “IfD method” ==> “LfD” method?
- Algorithm 1, page 6: what are the $\phi$ and $w$ parameters? It looks to me that they are taken from WGAN paper but with no explanation.
- Two references “Syed et al. (2008a)” and “Syed et al. (2008b)” are actually the same.
- Reference “Langlois et al. (2019)” should be corrected as “Wang et al. (2019)” -- see https://arxiv.org/abs/1907.02057