           Discount predictor: it p(it hit, zit)


where the RSSM model [1]} is used to learn the dynamics of the environment and Observation, Reward and Discount predictors are used for reconstructing corresponding variables.
We learn Observation, Reward and Discount predictors similar to [2]} via supervised loss and Representation model by minimizing the evidence lower bound [3]}. Transition predictor is then trained by minimizing KL divergence between \(\hat{z}^i_t\)  and \(z^i_t\) , where \(z^i_t\)  is a stochastic state sampled from categorical distribution same as in DreamerV2. Recurrent model consists of 1 layer GRU [4]} as in the original DreamerV2 implementation. In order to facilitate independence between hidden states \(\textbf {h}_t\) , we predict individual reward, discount and observation only from agent's latent state. We also use the sum of KL divergences between individual stochastic states \(\hat{z}^i_t\)  by applying the chain rule to the true Transition function:
\(p_\phi (\hat{\textbf {z}}_t \mid \textbf {h}_t) =p_\phi (\hat{z}^1_t \mid \textbf {h}_t) \cdots p_\phi (\hat{z}^n_t \mid \textbf {h}_t, \hat{z}^1_t, \ldots , \hat{z}^{n-1}_t)\) 
