Temporal Difference Variational Auto-Encoder

1 INTRODUCTION

• For partially observed reinforcement learning, agents need representations of uncertainty gathered from past observations to act and explore optimally.The introduction notes that weak reinforcement signals may make such representations difficult to learn implicitly through model-free reinforcement learning.Source: §1 · p. 1
• Planning requires imagining distant futures consistent with the past, while step-by-step planning may be cognitively or computationally unrealistic.Source: §1 · p. 1
• TD-VAE is proposed to combine abstract state representations, belief states, and temporally abstract predictions, and its full model is illustrated in complex 3D environments.The paper describes jumpy future predictions and filtering-posterior training as parts of this contribution.Source: §1 · pp. 1–2
• The new inference network and modified lower bound improve likelihood over classical deep state-space training methods on a simple task.The comparison is reported for a simple task; the passage does not name the individual baseline methods.Source: §1 · p. 2

2.1 CONSTRUCTION OF A LATENT STATE-SPACE

• Autoregressive models predict in observation space without a compressed representation, whereas latent state-space models support state-to-state sampling without rendering observations.Autoregressive test-time feedback can also be unstable because the model is trained on real next-step inputs but receives its own predictions at test time.Source: §2.1 · p. 2

2.2 ONLINE CREATION OF BELIEF STATE.

• Under the state-space model's Markov assumption, the filtering distribution over the current latent state given observations so far is sufficient to compute the conditional future distribution.A neural code approximating that filtering distribution can therefore serve as a belief state.Source: §2.2 · p. 3
• Classical autoregressive posterior training may leak uncertainty about the current latent state's marginal posterior into a sampled earlier latent state.Recovering all information may then require resampling the earlier latent states back to the beginning of the sequence.Source: §2.2 · p. 3
• Learned deep belief states remain comparatively underexplored; cited related approaches include neural particle filtering and predictive state representations.Source: §2.2 · p. 3

3 BELIEF-STATE-BASED ELBO FOR SEQUENTIAL TD-VAE

• For each time t, the model evaluates the observation likelihood by inferring over only z_(t−1) and z_t, yielding belief states at adjacent times.The derivation begins from an autoregressive decomposition of the data likelihood and uses the state-space model's Markov assumptions.Source: §3 · p. 3
• The model uses a recurrent code b_t=f(b_(t−1),x_t) to represent beliefs at both times and a one-step smoothing posterior for z_(t−1).The previous belief and current filtering posterior are approximated with the same belief distribution.Source: §3 · p. 4
• The resulting loss combines the observation likelihood, the preceding state's belief distribution, and the state-transition probability.Source: §3 · p. 4

4 TD-VAE AND JUMPY STATE MODELING

• Jumpy state predictions can support planning across different timescales, provide stronger future training signals, and avoid modeling intervals that do not affect later predictions.The authors also identify computational efficiency and independence from arbitrary temporal subsampling as motivations.Source: §4 · p. 4
• Unlike approaches that predict skipped observations, TD-VAE models transitions between distant latent states, preserving a state representation for recursive prediction in partially observed settings.The cited alternatives may discard skipped observations or target predictable frames without learning appropriate states.Source: §4 · p. 4, §4.1 · p. 5

4.1 THE TD-VAE MODEL

• A forward RNN encodes past observations into beliefs, and TD-VAE relates states at t1 and t2 with a jumpy transition and a loss that includes future reconstruction and an earlier belief prior.The model can sample t1 and a separation t2−t1 from a finite range; the paper also develops a hierarchical version with multiple latent-state layers.Source: §4.1 · p. 5

4.2 INTUITION BEHIND TD-VAE

• The loss tests whether a future state can be predicted from an earlier state by comparing the transition with a smoothed estimate of the earlier state conditioned on information available at both times.The smoothing distribution is also matched to the earlier belief distribution.Source: §4.2 · p. 6

4.3 CONNECTION WITH TEMPORAL-DIFFERENCE LEARNING

• TD-VAE applies a temporal-difference analogy by using a later belief to train an earlier belief through variationally auto-encoding a future-state sample across the two times.The state-to-state mapping is intended to make beliefs at the two time steps consistent.Source: §4.3 · p. 6

5 EXPERIMENTS.

• The first experiment uses sequential TD-VAE for comparison with related state-space model algorithms; subsequent experiments use the full model.Source: §5 · p. 6

5.1 PARTIALLY OBSERVED MINIPACMAN

• TD-VAE outperforms both baselines on MiniPacman test-set ELBO and estimated negative log probability, while the mean-field model performs worst.The comparison also reports that naively restricting the encoder to obtain a belief state hurts performance significantly, which TD-VAE overcomes.Source: §5.1 · pp. 6–7, Figure 2 · p. 7

5.2 MOVING MNIST

• Moving MNIST rollouts skip multiple time steps while preserving the digit’s direction of motion.The model samples a state from its belief and repeatedly transitions between states before decoding the frames.Source: §5.2 · p. 7, Figure 3 · p. 7

5.3 NOISY HARMONIC OSCILLATOR

• The oscillator model predicts frequency and magnitude at both dt = 20 and dt = 100, while phase is less accurate at dt = 100 because noise makes it unpredictable.A separate belief-state classifier experiment finds TD-VAE compares better with next-step prediction as the separating interval grows and the LSTM gets smaller.Source: §5.3 · p. 8, Figure 4 · p. 8

5.4 DEEPMIND LAB ENVIRONMENT

• In DeepMind Lab, samples from one belief decode to roughly the same frame, but their predicted futures differ, representing several possible futures.The experiment uses time skips sampled uniformly from [1, 40].Source: §5.4 · p. 9, Figure 5 · p. 9
• DeepMind Lab rollouts trained on time separations uniformly sampled from [1, 5] appear to move forward into corridors while skipping several time steps.The reported single-step motion is slower than the motion in these rollouts.Source: §5.4 · p. 9, Figure 6 · p. 9

6 CONCLUSIONS

• Random-interval bridges let TD-VAE states relate directly over longer stretches and explicitly encode the future.The model can roll out in state space using steps larger than, and potentially independent of, the underlying data step size.Source: §6 · p. 9
• TD-VAE is presented as a model different from an accurate step-by-step environment simulator.The authors identify applying it to more complex settings and investigating reinforcement-learning uses such as representation learning and planning as future work.Source: §6 · p. 9

A TD-VAE AS A MODEL OF JUMPY OBSERVATIONS

• The jumpy objective models a subsequence with latent states linked by transitions and observations generated from those states.The sequence may be sampled according to an arbitrary distribution over subsequences.Source: A TD-VAE AS A MODEL OF JUMPY… · p. 13
• Full-sequence belief states enrich the posterior over sampled latent states beyond dependence on the observed subsequence alone.Observations outside the sampled subsequence serve as auxiliary variables for the posterior.Source: A TD-VAE AS A MODEL OF JUMPY… · p. 13
• The resulting objective uses the same loss as one-step TD-VAE when the sampling scheme selects only pairs of time points.The derivation modifies the approximate ELBO for the jumpy training regime.Source: A TD-VAE AS A MODEL OF JUMPY… · p. 13

B DERIVATION OF THE TD-VAE MODEL FROM ITS DESIRED PROPERTIES

• A general recurrent variational auto-encoder uses forward and backward encoder RNNs with a forward decoder, while allowing other connections.The architecture can generate sequences one step at a time and can therefore generate arbitrarily long sequences.Source: B DERIVATION OF THE TD-VAE MODEL FROM… · p. 13
• Forward sampling in latent space requires the encoder not to feed into the decoder or latent prior, because encoding requires observations.Without connections 2, the recurrent VAE can sample forward in latent space without returning to observations.Source: B DERIVATION OF THE TD-VAE MODEL FROM… · p. 13, Figure 7 · p. 14
• The belief state represents knowledge up to the current time, so it cannot receive input from a backward decoder.The forward encoder is proposed for computing the belief state because it should have unrestricted access to information without sampling or bottlenecks.Source: B DERIVATION OF THE TD-VAE MODEL FROM… · pp. 13–14
• The model makes the prior and reconstruction depend only on the latent state, rather than also on decoder state.This choice avoids the possibility that decoder state stores information that would otherwise need to be represented by the encoder.Source: B DERIVATION OF THE TD-VAE MODEL FROM… · p. 14
• Smoothing conditions an earlier posterior on the future; the box example motivates sampling the later state first, then inferring the earlier one.Otherwise, an uncertain earlier belief could sample the wrong object and teach an incorrect transition.Source: B DERIVATION OF THE TD-VAE MODEL FROM… · p. 14
• Smoothing uses a backward encoder, while the state-transition model remains latent-to-latent.The backward encoder's access to future information rules out including decoder state in the belief state.Source: B DERIVATION OF THE TD-VAE MODEL FROM… · p. 14
• A KL term trains the prior belief distribution to stay close to the posterior, inducing training of the backward encoder through its posterior role.Source: B DERIVATION OF THE TD-VAE MODEL FROM… · p. 14
• For separated times, the algorithm samples the later belief as a posterior against a transition prior, reconstructs its observation, and infers the earlier state from the later sample.The earlier-state loss compares its posterior conditioned on both beliefs with its belief prior.Source: B DERIVATION OF THE TD-VAE MODEL FROM… · p. 15

C HIERARCHICAL MODEL

• The hierarchical model is intended to represent higher-level states that predict same-level or lower-level states and capture more invariant or abstract information.Source: C HIERARCHICAL MODEL · p. 15
• Its belief aggregator is a deep LSTM whose layers receive inputs from the next higher layer at the previous time step, allowing cross-layer influence.Source: C HIERARCHICAL MODEL · p. 15
• Sampling proceeds from higher-level representations toward lower layers, closer to pixels, and inference follows the same downward order.Source: C HIERARCHICAL MODEL · p. 15
• Each layer's posterior depends on samples from previously sampled layers, summarized by an RNN; weights are not shared across layers.Source: C HIERARCHICAL MODEL · p. 16

• For harmonic-oscillator experiments, belief states come from an LSTM, while context maps parameterize normal distributions and share weights when contexts differ only by time index.Source: D FUNCTIONAL FORMS AND PARAMETER CHOICES · p. 16
• D-map hidden layers and belief states have size 50, each z_l has size 8, and Adam uses learning rate 0.0005.Source: D FUNCTIONAL FORMS AND PARAMETER CHOICES · p. 16
• For MNIST, observations pass through a two-hidden-layer ReLU MLP, the decoder outputs Bernoulli logits, and δt is omitted from network inputs.Source: D FUNCTIONAL FORMS AND PARAMETER CHOICES · p. 16
• For DeepMind Lab, fully connected LSTMs have size 512, convolutional LSTMs use 4 × 4 × 256 layers with 3 × 3 kernels, and the decoder adds a DRAW-like canvas layer.Source: D FUNCTIONAL FORMS AND PARAMETER CHOICES · p. 17