Why does the RSSM receive different inputs for learning and imagination? about dreamer HOT 3 CLOSED

Yes, you can think of the prior as a loose guess about the next state and of the posterior as a more precise guess, because it is based on additional evidence, namely the observed image.

During training images are available and the GRU takes a posterior sample as input. During imagination, the future images are not available and the GRU takes a prior sample as input instead.

Intuitively, this works because the posterior and prior are trained to be close to each other using the KL regularizer. And because the random samples have some spread during training, the GRU becomes robust to seeing new samples during imagination.

Mathematically, you are right that this simply follows from the evidence lower bound.

When interacting with the environment, the current image is available, and so we use the posterior.

from dreamer.

Check out the PlaNet paper where we introduced this model, especially Figure 2c. You are right that the prior and posterior transitions differ. Specifically, the prior does not have access to the embedding of the image. This allows predicting ahead without knowing the images. In the figure I pointed you to, the solid lines are always used but the dashed lines are only used by the posterior. I don't think your diagrams are correct, or I don't quite understand them. I would draw them as

Prior:               Posterior: 
                       
a ---.               a ---. 
     |                    | 
     |                    | 
h ---+---> h'        h ---+---> h'  
     |     |              |     |  
     |     v              |     v  
s ---'     s'        s ---'     s'  
                                ^
                                |
                                o'

from dreamer.

Hi @danijar,

Sorry that I misplaced the some of notations in the diagrams and left the previous issue unorganized.

Here's a figure I draw about the RSSM

The left part shows how the RSSM looks like in a single time step when we call RSSM.observe(this line), where $e$ is the image embedding, $s'$ is the posterior state, and $s$ is the prior state. The right part shows how the RSSM looks like when we call dreamer._imagine_ahead(this line). My point, if I understand them right, is that the GRU takes $s'$ from the posterior as input when we call RSSM.observe. In contrast, when we call dreamer._imagine_ahead, the GRU takes $s$ from the prior as input. This different behavior confused me for a while. Then I realized that when we compute the KL divergence in the ELBO, we take the expectation under the posterior. Furthermore, the dynamics model is trained using the data collected by the policy, which is conditioned on the feature from the posterior. Both make sense of the left diagram.

Hope these can help you understand the figure I drew. Of course, if I made any mistakes, please let me know.

from dreamer.