Reparameterization trick for Q_psi

Killed 2 bugs with one stone - turns out that I don't need to compute $C_t=A\Sigma A^T+H$ in order to generate a sample from it!

To sample $z_{t+1} \sim N(A\mu+Bu+o, A\Sigma A^T+H)$, let $\omega\sim N(0,H)$ and $x \sim N(0,\Sigma)$. $\Sigma$ is diagonal so we can sample $x$ using the ordinary reparameterization trick outlined in the VAE paper.

Let $y=Ax+\omega$.

Then we have
$$\mu_y=E[Ax+\omega]=E[Ax]+E[\omega]=A*0+0=0$$
$$\Sigma_y=\text{Var}[Ax+\omega]=\text{Var}[Ax]+\text{Var}[\omega]=A\Sigma A^T + H$$

Which yields us the exact covariance we want!

Then we can just compute $z_{t+1}=A\mu+Bu+o+ y$. Note that we could have centered $x$ at $\mu$, in which $\mu_y=A\mu$, to which we just add $Bu+o$, but it turns out that we need to compute $\mu_{z_{t+1}}$ anyway for measuring the KL divergence later.