Predictive-State Decoders: Encoding the Future into Recurrent Networks
Recurrent neural networks (RNNs) are a vital modeling technique that rely on internal states learned indirectly by optimization of a supervised, unsupervised, or reinforcement training loss. RNNs are used to model dynamic processes that are characterized by underlying latent states whose form is often unknown, precluding its analytic representation inside an RNN. In the Predictive-State Representation (PSR) literature, latent state processes are modeled by an internal state representation that directly models the distribution of future observations, and most recent work in this area has relied on explicitly representing and targeting sufficient statistics of this probability distribution. We seek to combine the advantages of RNNs and PSRs by augmenting existing state-of-the-art recurrent neural networks with PREDICTIVE-STATE DECODERS (PSDs), which add supervision to the network's internal state representation to target predicting future observations. PSDs are simple to implement and easily incorporated into existing training pipelines via additional loss regularization. We demonstrate the effectiveness of PSDs with experimental results in three different domains: probabilistic filtering, Imitation Learning, and Reinforcement Learning. In each, our method improves statistical performance of state-of-the-art recurrent baselines and does so with fewer iterations and less data.
The authors present a simple method for regularizing recurrent neural networks (specifically applied to low-dimensional control tasks trained by filtering, imitation, and reinforcement learning) that is inspired by the literature on Predictive State Representations (PSRs). In PSRs, processes with latent states are modeled by directly identifying the latent state with a representation of the sufficient statistics of future observations. This stands in contrast to RNNs, which use an opaque hidden state that is optimized by backpropagation to make the whole system model the target of interest through evolution of the hidden state. They propose a simple regularization method for RNNs inspired by PSRs, called Predictive State Decoders, which maintains the parametric form of a standard RNN, but augments the training objective to make the hidden state more directly predict statistics of distant future observations. While direct training of an RNN by backpropagation implicitly encourages the internal state to be predictive of the future, the authors demonstrate that this explicit additional regularization term gives significant improvements on several control tasks.
The idea specifically is to augment the standard training loss for an RNN controller (either for filtering, imitation learning, or reinforcement learning) with an additional term that maps the hidden state at each timestep through an additional learned parametric function F, and penalizes the difference between that and some statistics of up to 10 future observations. This encourages the internal state of the RNN to more directly capture the future distribution of observations without over-fitting to next-step prediction, and architecturally improves gradient flow through the network similar to skip-connections, attention, or the intermediate classifiers added in FitNets or Inception.
I like the main idea of the paper. It is simple, well-motivated, and connects RNN-based models to a large body of theoretically-grounded research in modeling latent state processes, of which the authors do a good literature review. The experiments and analysis are thorough, and demonstrate significant improvements over the baseline on many of the tasks in the OpenAI gym. The baselines use strong, well-known techniques and optimization methods for imitation learning and RL.
I have some concerns that the paper does not offer enough detail about the exact architectures used. For the filtering experiment, it is stated that phi is the identity function, which I assume is true for the future experiments. Additionally, unless I am missing something, the exact functional form of F, the decoder, is never mentioned in the text. Is this a feedforward neural network? Given that the predicted states can sometimes be 10 times the size of the observation dimension, have the authors thought about how to apply this method to models with very large dimensional observations? More detail would be required on exactly the form of F in a camera-ready.
Overall, I think this is a well-executed paper that demonstrates a nice synergy between the PSR and RNN methods for latent state process modeling. The method seems simple to implement and convincingly improves over the baselines.