Question: Is the reward function learned as well?
Evidence: Updating the policy. The policy is optimized to maximize the augmented reward on transitions sampled from the learned dynamics model. While this optimization can be done using any RL algorithm, including on-policy methods, we will focus on an off-policy actor-critic method.
We define the Q function as sum of augmented rewards under the learned dynamics model:
We approximate the Q function using a neural network Q ψ (s t , a t ) with parameters φ. We train the Q function using the TD loss on transitions sampled from the learned dynamics model:
where • sg is the stop-gradient operator and
) . The augmented reward, r, is estimated using the learned classifier. To estimate the corresponding value function, we use a 1-sample approximation:
The policy is trained to maximize the expected future (augmented) return, as estimated by the Q function:
In our implementation, we regularize the policy by adding an additional entropy regularizer. Following prior work, we maintain two Q functions and two target Q functions, use the minimum of the two target Q functions to compute the TD target. See Appendix C for details.
Updating the classifier. We train the classifier to distinguish real versus model transitions using the standard cross entropy loss: between training the policy on experience from the learned dynamics model with augmented rewards and updating the model+classifier using a GAN-like loss. Updates are gradient steps with the Adam optimizer.
1: while not converged do 2:
Sample experience from learned model and modify rewards using the classifier (Eq. 5).
Update policy and Q function using the model experience and augmented rewards (Eq.s 8 and 7). 4:
Update model and classifier using GAN-like losses (Eq.s 9 and 10). 5:
(Infrequently) Sample experience from real model. 6: return policy π θ (at | st).
Updating the dynamics model. To optimize the dynamics model, we rewrite the lower bound in terms of a single transition (derivation in Appendix A.6):
The approximation above reflects approximation error in learning the optimal classifier. This approximation is standard in prior work on GANs and adversarial inference. The procedure for optimizing the dynamics model and the classifier resembles a GAN: the classifier is optimized to distinguish real transitions from model transitions, and the model is updated to fool the classifier (and increase rewards). However, our method is not equivalent to simply replacing a maximum likelihood model with a GAN model. Indeed, such an approach would not optimize a lower bound on expected return. Rather, our model objective includes an additional value term and our policy objective includes an additional classifier term. These changes enable the model and policy to optimize the same objective, which is a lower bound on expected return.
Algorithm summary. We summarize the method in Alg. 1 and provide an illustration in Fig.. Implementing MnM on top of a standard model-based RL algorithm is straightforward. First, create an additional classifier network. Second, instead of using the maximum likelihood objective to train the model, use the GAN-like objective in Eq. 10 to update both the model and the classifier. Third, add the classifier's logits to the predicted rewards (Eq. 5). Following prior work, we learn an additional model to predict the true environment.
Algorithm summary. We summarize the method in Alg. 1 and provide an illustration in Fig. 1. Implementing MnM on top of a standard model-based RL algorithm is straightforward. First, create an additional classifier network. Second, instead of using the maximum likelihood objective to train the model, use the GAN-like objective in Eq. 10 to update both the model and the classifier. Third, add the classifier’s logits to the predicted rewards (Eq. 5). Following prior work (Janner et al., 2019), we learn a neural network to predict the true environment rewards.
Answer: Yes, the reward function is learned as well. We fit a separate neural network to predict the rewards (final paragraph of Section 4.2).