Episode Transformer: Model-based Episodic Reinforcement Learning

• The writing of this paper is hard to follow. For example, the results in Figure 3 and Figure 4 can be placed in one large figure, as they demonstrates similar results in different tasks, so that the readers can compare the performance more easily.

• As shown in Figure 6, the state prediction technique seems useless. However, this technique is core in model-based RL.

• It is rather straightforward to use the idea of learning a world model with a Transformer. The novelty lies in the Blackbox Reinforcement Learning setting, which introduced in prior works
• The experiment result is not good. For example, in Fig. 3a, the proposed method performs worse than others. Overall, in term of reward performance, the final convergence of the method is not better than simply using trust region optimization. This weakens the main message of using Episode Transformer. There are some improvements in terms of sample efficiency. However, that comes without surprise due to using the world model
• The ablation study is not convincing. Fig. 6 (middle and right) shows no clear difference.

The experiments section compares using only non-standard tasks, so it is somewhat difficult to tell whether the baselines are tuned appropriately. Specifically, the 5D reacher task seems quite simple (given that it is still a planar reaching problem) and I would expect most current RL methods to solve it to a similar level of performance under reasonable hyperparameter settings. The authors should add at least one "standard" environment to ensure fair comparison. It would be also good to compare against a modern sample-efficient RL method (e.g. REDQ, Dreamer, etc.) to validate the sample efficiency claims.

For the maximum jump height task, it may be helpful to include a baseline that keeps "max height achieved so far" as an additional state variable and applies SAC/PPO with the correct task reward (i.e. reward equal to $\max(0, h_t-\max_{1\le t \le t-1}h_t)$).

The ablations claim that ET works better with the stochastic backpropagation scheme than with REINFORCE, but the experiments don't seem to lend strong support for this and there seems to be quite a bit of overlap between the reported confidence intervals for the two. The authors should remove these claims or run a new experiment to tighten the confidence intervals and confirm whether there is some significant effect.

Given the ablations, it's not clear why it's necessary to predict true environment states in the first place. It seems that it would simplify the presentation of the method with no drawbacks if the state-prediction heads were removed. Additionally, as far as I can tell no ablation compares ET against a simpler model. This would be helpful to justify the Transformer model (e.g. can the same thing be accomplished with a smaller reward model?)

The introduction and related work should be reworked. Particularly, the explanation of Markovian rewards contains several questionable statements; e.g. the task of "throwing objects into bins" could absolutely be expressed with a Markovian reward - the problem in that case is with long-term credit assignment, in which case episodic RL could certainly perform better.

• How is the trajectory generator $\psi$ learned? Is the $\psi$ trained the same way for BB-TRPL and BB-PPO? How well is it trained?
• About the Trust Region Constraints (TRC):
  • I don't entirely agree with the statement, especially "In the step-based case, smaller errors during action selection can still be corrected at a later time step." In model-based RL, since the policy is learned during imagination, smaller errors are likely to be accumulated along rollouts instead of corrected at a later time step. This is known as the problem of compounding error.
  • I also doubt the argument behind TRC since the policy is trained with on-policy data during imagination, the learning policy is always in the "trust region". Moreover, equation 6 optimizes the expected return with respect to $\tilde{\pi}(w | c, \theta)$ instead of the learning policy. The policy is learned indirectly via the KL regularizer. I wonder if the learning is effective in this way.
  • The moving average of policy weights here reminds me of other actor-critic methods, so I'm curious about that, instead of using the TRC, maintaining a moving average of $R_\phi(c, w)$, directly optimizing the learning policy via $E_\pi(w | c)[- \tilde{R}_\phi(c, w)]$ might also help.
• Could the authors prove the necessity of using a Transformer by comparing the proposed method with an RNN-based world model? Especially when predicting the auxiliary state is not very important, I'm curious why not using a more lightweight RNN-based world model.