Sparsely Changing Latent States for Prediction and Planning in Partially Observable Domains
Authors: Christian Gumbsch
Martin V. Butz
Georg Martius
Venue: NeurIPS 2021 Poster (2021)
Keywords: recurrent neural networks
partial observability
sparsity
explainability
generalization
model-based planning
reinforcement learning

Abstract: A common approach to prediction and planning in partially observable domains is to use recurrent neural networks (RNNs), which ideally develop and maintain a latent memory about hidden, task-relevant factors. We hypothesize that many of these hidden factors in the physical world are constant over time, changing only sparsely. To study this hypothesis, we propose Gated $L_0$ Regularized Dynamics (GateL0RD), a novel recurrent architecture that incorporates the inductive bias to maintain stable, sparsely changing latent states.  The bias is implemented by means of a novel internal gating function and a penalty on the $L_0$ norm of latent state changes. We demonstrate that GateL0RD can compete with or outperform state-of-the-art RNNs in a variety of partially observable prediction and control tasks. GateL0RD tends to encode the underlying generative factors of the environment, ignores spurious temporal dependencies, and generalizes better, improving sampling efficiency and overall performance in model-based planning and reinforcement learning tasks. Moreover, we show that the developing latent states can be easily interpreted, which is a step towards better explainability in RNNs.

Sparsely Changing Latent States for Prediction and Planning in Partially Observable Domains
Christian Gumbsch
Martin V. Butz
Autonomous Learning Group
Neuro-Cognitive Modeling Group University of Tübingen
Max Planck Institute for Intelligent Systems
& Neuro- Neuro-Cognitive Modeling Group University of Tübingen Tübingen, Germany
Tübingen, Germany
martin. butz@uni Suttzunni-tubbingen.de
christian, www.aaanennnneeeeennnnnn........ mpg.de
Georg Martius
Autonomous Learning Group
Max Planck Institute for Intelligent Systems
Tübingen, Germany
georg goorg,.artiiuutuuubme.e.....
Abstract
A common approach to prediction and planning in partially observable domains is to use recurrent neural networks (RNNs), which ideally develop and maintain a latent memory about hidden, task-relevant factors. We hypothesize that many these hidden factors in the physical world are constant over time, changing only sparsely. To study this hypothesis, we propose Gated Ly Regularized Dynamics (GateLORD), novel recurrent architecture that incorporates the inductive bias to maintain stable, sparsely changing latent states. The bias is implemented by means of novel internal gating function and penalty the Lo norm of latent state changes. We demonstrate that GateLORD can compete with or outperform state-of-ther -hee--rrtte-rttt RNNs in variety of partially observable prediction and control tasks. GoteLORD tends to encode the underlying generative factors of the environ- ment, ignores spurious temporal dependencies, and generalizes better, improving sampling efficiency and overall performance in model-based planning and rein- forcement learning tasks. Moreover, we show that the developing latent states be easily interpreted, which a step towards better explainability in RNNs.
1 Introduction
When does the meeting start? Where are my car keys? Is the stove turned off? Humans memorize lots of information over extended periods of time. In contrast, classical planning methods assume that the state of the environment fully observable at every time step [1]. This assumption does not hold for realistic applications, where generative processes are only indirectly observable entities are occluded. Planning in such Partially Observable Markov Decision Processes (POMDP) is challenging problem, because quitably-structurg memory is required for decision making.
Recurrent neural networks (RNNs) are often used to deal with partial observability [2-4]. They encode past observations by maintaining latent states, which iteratively updated. However, continuously updating the latent state causes past information to quickly "wash out". Long-Sho Term Memory networks (LSTM, [5]) and Gated Recurrent Units (GRU, [6]) deal with this problem by using internal gates. However, they cannot leave their latent states completely unchanged, because
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

small amounts of information continuously leak through the sigmoidal gating functions. Additionally. inputs typically need to pass through the latent state affect the output, making it hard to disentangle observable from unobservable information within their latent states.
Our hypothesis is that many generative latent factors in the physical world are constant over extended periods of time. Thus, there might not be the need to update memory at every time step. For example, consider dropping an object: If the drop-off point well some latent generative factors, such as gravity and aerodynamic object properties, are known, iteratively predicting the fall can be reasonably accomplished by a non-recurrent process. Similarly, when an agent picks up a key, is sufficient to memorize that the key is inside their pocket. However, latent factors typically do change significantly and systematically at particular points in time. For example, the aerodynamic properties of an object change drastically when the falling object shatters on the floor, and the location of the key changes systematically when the agent removes from their pocket.
These observations are related to assumptions used in causality research. A common assumption is that the generative process of system composed of autonomous mechanisms that describe causal relationships between the system's 's variables [7-9]. When considering Markov Decision Processes, it has been proposed that these mechanisms tend to interact sparsely in time and locally in space [10, 111. Causal models aim at creating dependencies between variables only when there exists a causal relationship between them, in order to improve generalization [8]. Updating the latent state of a model in every time step, on the other hand, induces the prior assumption that the generative latent state typically depends on all previous inputs. Thus, by suitably segmenting the dependencies of the latent variables over time, one can expect improved generalization across spurious temporal dependencies.
Very similar propositions have been made for human cogmition. Humans tend to perceive their stream of sensory information in terms of events [12-16]. Event Segmentation Theory (EST) [16] postulates a set of active event models, which encode event-respective aspects over extended periods of time and switch individually at event transitions. To learn about the transitions and consolidate associated latent event encodings, measurements of surprise and other significant changes in predictive model activities, as well as latent state stability assumptions, have been proposed as suitable inductive event segmentation biases [16-22]. Explicit relations to causality have been put forward in [23].
In accordance to EST and our sparsely changing latent factor assumption, we introduce Gated Lo Regularized Dynamics (GateLORDD). Gatel.ORD applies Lo-regularized gates, inducing an inductive learning bias to encode piecewise constant latent state dynamics. GateLORD thus becomes able to memorize task-relevant information over long periods of time. The main contributions of this work can be summarized as follows. (i) We introduce a stochastic, rectified gating function for controlling latent state updates, which we regularize towards sparse updates using the Lo norm. (ii) We demonstrate that our network performs as good or better than state-of- oftte-of-the-arr art RNNs for prediction or control in various partially- -owaalll--obeevbblee problems with piecewise constant dynamics. (iii) We also show that the inductive bias leads to better generalization under distributional shifts. (iv) Lastly, we show that the latent states can be easily interpreted by humans.
2 Background
Let fa : X x H be recurrent neural network (RNN) with learnable parameters 0 mapping inputs' x; E X and h E H the latent (hidden) state to the output ÿ. E y and updated latent states h. The training dataset D consists of sequences of input-output pairs d [(x1, Vy),.... , (ær. yr)) of length T. In this paper, we consider the prediction and control of systems that can be described by partially observable Markov decision process (POMDP) with state space S, action space A, observations space O, and deterministic hidden transitions x A S.2
Notation: bold lowercase letters denote vectors (e.g.. æ). Vector dimensions are denoted by superscript (e.g. 2 Jar", x²? ,,,...,""" _1 E R.). Time or other additional information denoted by subscript (c.g., ..)) We treat the prediction of time series without any actions special case of the POMDP with A 0.

---

3 Le-regulariztionn of latent state changes
We want the RNN fo to learn to solve a task, while maintaining piecewise constant latent states over time. The network creates dynamics of latent states h, when applied to a sequence: (Øt, hx) So (zt; 1 starting from some fg. The most suitable measure to determine how much a time-series piecewise constant is the Lg norm applied temporal changes. With the change in latent state as A h4 = h hi, we define the Ly-loss as
CL (Ah) = | = Σ I( #0),
(1)
J=1
which penalizes the number of non-zero entries of the vector of latent state changes Ah.
The regularization loss from Eq. I can combine in the usual way with the task objective to yield the overall learning objective L of the network:
((D,0) = EAADD Σ Craste( (Dr, ye) + A (Ahi)
(2)
with (ū., = (a 11-1). The task-dependent loss Lrask (+: can be, for instance, the mean-
squared error for regression or cross-entropy loss for classification. The hyperparameter À controls the trade-off between the task-based loss and the desired latent state regularization.
Unfortunately, we cannot directly minimize this loss using gradient- based techniques, such as stochastic gradient descent (SGD), due to the non-differentiabiiiy of the Lo-term. Louizos et al. [24] proposed a way to learn Ly regularization of the learnable parameters of a neural network with SGD, They achieve this by using set of stochastic gates controlling the parameters' usage. Each learnable parameter 6 that is subject to the Lg loss is substituted by a gated version 8 O(s))) ) where (-)) is the Hcaviside step function (+(s) 0 if S 0 and otherwise) and is determined by distribution q(() v) with learned parameters v. Thus, O' is only non-zero s > 0. This allows to rewrite the Lo loss (Eq. 1) for 0' as:
Cr. (0', = ==0 Σ 3
with ~ (()
where parameters v influence sparsity and affected by the loss.
To tackle the problem of non-differentiable binary gates, we can use a smooth approximation as a surrogate [24-26]. Alternatively, we can substitute its gradients during the backward pass, for example using the straight-through estimator [27]. which treats the step function as a linear function during the backward pass, or approximate its gradients as in the REINFORCH algorithm [28].
To transfer this approach to regularize the latent state dynamics in an RNN, we require an internal gating function A(-) E [0, 1]. which controls whether the latent state is updated or not. For instance:
h h 1 A(s)
with A h h
where h is the proposed new latent state and is stochastic variable depending on the current input and previous latent state and the parameters, i.e. 8, ~ q(8;; I,, 1,,,,,1)., vv.. J. For brevity, we merge the parameters into the overall parameter set, i.e. O. For computing Eq. 2 we need binarize the gate by applying the step function ((A(s)). Thus we can rewrite Eq. 2 as
0) -END ELosk Craak (#r. y1) + [O(A(s:))
LSTMs and GRUs use deterministic sigmoidal gates for A in Eq. 4 to determine how to update their latent state. However, it is not straight forward to apply this approach to them (detailed in Suppl. A). Thus, we instead introduce novel RNN, that merges components from GRUs and LSTMs, to implement the proposed Lo regularization latent state changes while still allowing the network to make powerful computations. We name network Gated Lo Regularized Dynamics (GateLORD).
3

---

Ais)
post
E
ho
h
h
h
(b) sate-activation A(s)
h
fint
frm
Forward
20
Badiooor
(a) Illustration of the core of GateLORD. Figure I: Architecture overview. (a) GateLORD with its three subnetworks. The gating function controls the latent state update (red), the recommendation function computes a new latent state (blue) and the output function computes the output (purple). (b) Gate-activation function A (ReTanh). (c) Heaviside step function O and its gradient estimator. (d) Overall architecture.
O with its substitution
(d) Overall GateLORD
GateLORD
The core of GateLORD implements the general mapping hg) fx h1- using three func- tions, or subnetworks: (1) a recommendation network which proposes new candidate latent state, (2) gating network g. which determines how the latent state updated, and (3) an output
function, which computes the output based on the updated latent state and the input. The network is systematically illustrated in Fig. la.
The overall processing is described by the following equations:
St (g((,, 11
(sample gate input)
(6)
A(s) == max(0, tanh(s))
(new gating function)
(7)
h, h, + (sxe @ (r(#,, h(-1) h
(update keep latent state)
(8)
G. p(x,, h4) @ o(z,, h.),
(compute output)
(9)
where O denotes demennt-wis multiplication (Hadamard product).
We start with the control of the latent state in Eq. 8. Following Eq. 4, new latent value is proposed
by the recommendation function r(æ,, h1-1) and the update is "gated" by A(s). Importantly, if A(s) = 0 no change to the latent state occurs. Note that the update Eq. is in principle equivalent to the latent state update in GRUs [6]. for which it is typically written h, A(s) @ r(x,, hz-1)
(1 A(s)) 3 h;-1 with ((s) deterministic sigmoidal gate.
Because we aim for piecewise constant latent states, the gating function A defined in Eq. 7 needs to be able to output exactly zero. A potential choice would be the Heaviside function, i.e. either copy the new latent state or keep the old one. This, however, does not allow any multiplicative computation. So a natural choice is to combine the standard sigmoid gate of RNNs with the step- function: A(s) max(0, tanh(s) (s)) which we call ReTann (rectified tanh) Figure 1 shows the activation function depending on its input. The gate closed (A(s*) = 0) for all inputs s² S 0. A closed gate results in latent state that remains constant in dimension i, i.e., hj hi On the other hand, for s the latent state is interpolated between the proposed new value and the old one.
The next puzzle piece is the input to the gate. Motivated from the Lo regularization in Eq. I use stochastic input. However, in our RNN setting, it should depend on the current situation. Thus, we use Gaussian distribution for q with the mean determined by the gating network 0(x;,,,-) as defined in Eq. 6. We chose fixed diagonal covariance matrix Σ, which we set to Σ 0.1. To train our network using backpropagation, we implement the sampling using the reparametrization trick [29]. We introduce a noise variable and compute the gate activation as
8; g(,,, h4-1) +E
with 6 ~ NO, E).
(10)
During testing we set to achieve maximally accurate predictions.
Note that tanh(s) sigmoid(2s) 1.

---

Finally the output ÿ is computed from the inputs and the new latent state h in Eq. 9. Inspired by LSTMs [5], the output determined by a multiplication of normal branch (p(æ,, h:)) and a sigmoidal gating branch (o(æ;, hz)). We thus enable both additive as well as multiplicative effects of at and h2 on the output, enhancing the expressive power of the piecewise constant latent states.
In our implementation, all subnetworks are MLPs. T, use tanh output activation; o uses a sigmoid; g has linear output. p. o are one-layer networks. By default, r, g are also one-layer networks. However, when comparing against deep (stacked) RNNs, we increase the number of layers of r and to up to three (cf. Suppl. B).
We use the loss defined in Eq. 5. GateLORD is fully differentiable except for the Heaviside step function θ in Eq. 5. A simple approach to deal with discrete variables is to approximate the gradients by a differentiable estimator [25-27]. We employ the straight-through estimator [27], which substitutes the gradients of the step function θ by the derivative of the linear function (see Fig. 1c).
We use GoteLLRDD as a memory module of more general architecture illustrated in Fig. 1d. The network input is preprocessed by a feed-forward network fzr) Similarly, its output is postpro- cessed by an MLP fpoer( (i.e. a readout layer) before computing the loss. The latent state he of GateLORD could be initialized by 0. However, improvements can be achieved if the latent state is instead initialized by a context network fint, a shallow MLP that sets hg based on the first input [30, 31].
In the Supplementary Material we ablate various components of GateL.ORD, such as the gate activation function A (Suppl. C.1), the gate stochasticity (Suppl. C.2), the context network finit (Suppl. C.3), the multiplicative output branch o (Suppl. C.4), and compare against I=/T.g-variants (Suppl. C.5).
5 Related Work
Structural regularization of latent updates: Pioneering work on regularizing latent updates was done by Schmidhuber [32] who proposed the Neural History Compressor, a hierarchy of RNNs that autoregressively predict their next inputs. Thereby, the higher level RNN only becomes active and updates its latent states, if the lower level RNN fails to predict the next input. To structure latent state updates, the Clockwork RNN [33] partitions the hidden neurons of an RNN into separate modules, where each module operates at its own predefined frequency. Along similar lines, Phased LSTMs [34] use gates that open periodically. The update frequency in Clockwork RNNs and Phased LSTMs does not depend on the world state, but only on a predefined time scale.
Loss-based regularization of latent updates: For latent state regularization, Krueger and Memi- sevic [35] have proposed using an auxiliary loss term that punishes the change in Lz-norms of the latent state, which results in piecewise constant norms but not dynamics of the hidden states.
Binarized update gates: Closely related to our ReTanh, Skip RNNs [36] use binary gate to determine latent state update decisions. Similarly, Gumbel-Gate LSTMs [37] replace sigmoid input and forget gates with stochastic, binary gates, approximated by Gumbel-Softmax estimator [26]. Selective-Activatio with deterministic, binary gate and also incentivize sparsity. However, for GRUs the network output corresponds to the networks' latent state, thus, a piecewise constant latent state will result in piecewise constant outputs. All of these models were designed for classification or language processing tasks none were applied for prediction or control in a POMDP setup, which we consider here.
RNNs (SA-RNNs) [38] modify a GRU by masking the latent state
Attention-based latent state updates: Sparse latent state updates can also be achieved using at- tention [39-41]. Ncural Turing Machines [39] use an attention mechanism to update an external memory block. Thereby, the attention mechanism can focus and only modify a particular locations within the memory. Recurrent Independent Mechanisms (RIMs) [42] use a set of recurrent cells that only sparsely interact with the environment and one another through competition and bottleneck of attention. Recent extensions explore the update of the cells and the attention parameters at different time scales [43]. For RIMs the sparsity of the latent state changes is predefined via a hyperparameter that sets the number of active cells. In contrast, our Lo loss implements a soft constraint.
Transformers [41] omit memory altogether, processing complete sequence for every output at once using key-based attention. While this avoids problems arising from maintaining a latent state, their self-attention mechanism comes with high computational costs. Transformers
Transformers:
5

---

(a) Robot Remote Control
(b) Shepherd
Pick&Place
(d) MiniGrid
Figure 2: Simulations used to test GateLORD. (a) and (b) are continuous 2D-control tasks: (a) requires
triggering the control of a robot by getting remote control; (b) needs memorization of the sheep's position to capture later. (c) is the Fetch Pick& Place environment [47] modified to hecome partially observable and (d) shows problem (DoorKey-B) of the unni-Gridooldd suite [48].
have shown breakthrough success in natural language processing. However, it remains challenging to train them for planning or reinforcement learning applications in partiall--obeerablee domains [44].
Experiments
Our experiments offer answers the following questions: (a) Does GateLORD generalize better to ontoof-distribution inputs in partially observable domains than other commonly used RNNs? (b)
Is GateLORD suitable for control problems that require (long-term) memorization of information? (c) Are the developing latent states GateLORD easily interpretable by humans? Accordingly, we demonstrate both GateLORD's ability to generalize from 1-step prediction regime to autoregressive N-step prediction (Sec. 6.1) and its prediction robustness when facing action rollouts from different policies (Sec. 6.2). We then reveal precise memorization abilities (Sec. 6.3) and show that GateLORD is more sample efficient in various decision-making problems requiring memory (Sec. 6.4). Finally, examine exemplary latent state codes demonstrating their explainability (Sec. 6.5).
In our experiments we compare GateLORD to LSTMs [5], GRUs [6], and Elman RNNs [45]. We use the architecture shown in Fig. 1d for all networks, only replacing the core fe. We examine the RNNs both as a model for moeec-preeiitive control (MPC) as well as a memory module in a reinforcement learning (RL) setup. When used for prediction, the networks received the input 2; (o,, a.) with observations o₂ E O and actions a, E AA at time t and were trained to predict the change in observation, i.e. y: Δo1+1 (detailed in Suppl. B. 1). During testing the next observational inputs were generated autoregressively as 6 = o, + ÿ.- In the RL setting, the networks received as an input æ; O, the observation O, O and were trained policy and value estimations (detailed in Suppl. B.6). The networks were trained using Adam [46], with learning rates and layer numbers determined via grid search for each network type individually (cf. Suppl. B).
an actor-critic architecture to produce both
We evaluate GateLORD in variety of partially observable scenarios. In the Billiard Ball scenario single ball, simulated in realistic physics simulator, is shot on a pool table with low friction from a random position in a random direction with randomly selected velocity. The time series contain only the positions of the ball. This is the only considered scenario without actions.
Robot Remote Control is a continuous control problem where agent moves according to the two-dimensional actions a, (Fig. 2a). Once the agent reaches fixed position (terminal), a robot in another room is also controlled by the actions. The observable state is composed of the agent's position and the robot's 's position. Thus, whether the robot is controlled or not is not observable directly. When planning, the goal is move the robot to particular goal position (orange square).
Shepherd a challenging continuous control problem that requires long-term memorization (Fig. 2b). Here, the agent's actions a, are the two movement directions and grasp action controlling whether to pick up or drop the cage. The sheep starts at the top of the scene moving downwards with fixed randomly generated velocity. The sheep is then occluded by the wall, which masks its position from the observation. If the agent reaches the lever, the gate inside the wall opens, and the sheep appears again at the same horizontal position at the open gate. The goal is to get the sheep to enter the previously placed cage. The challenge to memorize the sheep's horizontal position exactly over potentially long time to place the cage properly and to then activate the lever during mental

---

Texiing prodiction
(b)
Toeting posdittion serce
Testing prodiction
(d)
Number gole opeaiigs
10-1
0.83
10
10
LSTM
Git
Fiman INN
0.001
12000
LS1M
.... 0.01
0.01
GateltiD 0.001
GabeLottD
0.1
MSS
10-2
10m>
10-12
10
5k
Figure 3: Billiard Ball results: prediction errors when trained using teacher forcing (a), or using scheduled sampling (b). GateLORD's prediction error (c) and mean number of gate openings (latent state updates) (d) for different values of A. Shaded areas show + one standard deviation.
simulation. The seven-dimensional observation Or is composed of the height of the occluder and the positions of all entities.
Fetch Pick&Place (OpenAI Gym v1, [47]) is benchmark RL task where a robotic manipulator has to move a randomly placed box (Fig. 2c). In our modified setting! the observable state is composed of the gripper- and object position and the relative positions of object and fingers with respect the gripper. The four-dimensionll actions control the gripper position and the opening of the fingers.
MiniGrid [48] a gridworld suite with a variety of partially observable RL problems. At every time f, the agent (red triangle in Fig. 2d) receives an image-like, restricted, ego-centric view (grey area) as its observation o (7 X 3-dimensional). can either move forward, turn left, turn right, or interact with objects via its oe--hot-encoder actions The problems vary largely in their difficulty, typically contain only sparse rewards, and often involve memorization, e.g., remembering that the agent picked up key. Suppl. B.7 details all examined MiniGrid environments.
6.1 Learning autoregressive predictions
First, we consider the problem of autoregressive N-step prediction in the Billiard Ball scenario. Here, during testing the networks receive the first two ball positions as input and predict a sequence of 50 ball positions. We first train the RNNs using teacher forcing, whereby the real inputs are fed to the networks. Figure 3a shows the prediction error for autoregressive predictions. Only GateLORD with latent state regularization (A = 0.001) is able to achieve reasonable predictions in this setup. The other RNNs seem to learn to continuously update their estimates of the ball's velocity based on the real inputs. Because GateLORD punishes continuous latent state updates, learning leads to updates of the estimated velocity only when required, i.e. upon collisions, improving its prediction robustness.
The problems of RNNs learning autoregressive prediction are well known [49, 50). A simple countermeasure is scheduled sampling [49], where each input is stochastically determined to be either the last network's 's output or the real input. The probability of using the network output increases over time. While the prediction accuracy of all RNNs improves when trained using scheduled sampling, GateLORD (A 0.001) still achieves the lowest mean prediction error (see Fig. 3b).
How does the regularization affect GateLORD? Figure 3c shows the prediction error for GateLORD for different settings of X. While a small regularization (A 0.001) leads to the highest accuracy in this scenario, similar predictions are obtained for different strengths (A E [0, 0.01)). Overly strong regularization (A- 0.1) degrades performance. Figure 3d shows the average gate openings per sequence. As indented, X directly affects how often GateLORD's latent state is updated: a higher value results in fewer gate openings and, thus, fewer latent state changes. Note that even for X == 0 GoteLORD learns to use fewer gates over time. We describe this effect in more detail in Suppl. D.I.
6.2 Generalization across policies
Particularly when priorities change or an agent switches behavior, different spurious temporal corrlattons can occur in the resulting sensorimotor timeseries data. Consequently, models are needed
We omit all velocities and the rotation of the object to make the scenario partially observable.
7

---

Teeting prefittioon
(b)
Gommibaauun Iene
(C)
Plarning
agens
Elsen INNN
lleas BAN
GRU
GRG
Lsua
0.6
Guoc LRD
QioloRD
MBE
0.6
Kiman RNNS
0.2
GRB
LSTM
10
LOOLD
25
06
84
10k
66
10k
20 40
spodia
gpode
Figure 4: Robot Remote Control results: prediction error on the test set (a) and on the generalization set (b). Success rate for MPC (c). Shaded areas show standard deviation (a & b) or standard error (c). Exemplary generalization sequence (d) showing the agent's ' positions (top), the robot's positions (middle) with GateLORD's 's predictions shown as dots, and GateLORD's latent states (bottom).
that generalize across those correlations. We use the networks trained as predictive models for the Robot Remote Control scenario investigate this aspect.
In Robot Remote Control the training data is generated by performing rollouts with 50 time steps of policy that produces random but linearly magnitue-inccrasinng actions. The actions' magnitude in the training data is positively correlated with time, which a spurious correlation that does not alter the underlying transition function of the environment in any way. We train the networks to predict the sequence of observations given the initial observation and sequence of actions. Thereby, test the networks using data generated by the same policy (test set) and generated by policy that samples uniformly random actions (generalization sef). Additionally, we use the trained RNNs for
control (MPC) using iCEM [, random shooting method that iteratively optimizes its actions to move the robot to the given goal position.
As shown in Fig. 4a, GateLORD (A 0.001) outperforms all other RNNs the test set. When tested on the generalization data, the prediction errors of the GRU and LSTM networks even increase over the course of training. Only GateLORD is able to maintain a low prediction error. Figure 4c shows the MPC performance. GateLORD yields the highest success rate.
Note that the lack of generalization is not primarily caused by the choice of hyperparameters even when the learning rate of the other RNNs was optimized for the generalization set, GateLORRD still outperformed them (additional experiment in Suppl. D.3). Instead, GoteLs better performance likely because it mostly encodes unobservable information within its latent state h, This is shown exemplarily in Fig. 4d (bottom row) and analyzed further in Suppl. D.5. The latent state remains constant and only one dimension changes once the agent controls the robot's position (middle row) through its actions. Because the other RNNs also encode observable information, e.g. actions, within their latent state, they are more negatively affected by distributional shifts and spurious dependencies.
GateLORD's improved generalization across temporal dependencies also holds for more compli- cated environments. In an additional experiment in Suppl. D.7 we show similar effects for the Fetch Pick PPace environment when trained on generalize across timings of the grasp.
sequences and tested to
(a)
Testing prediction erree
(b)
Shoep reoppearance ettur
(c)
Plancing stccess
10
10-1
Eluan RNN
MSSE of Ayheep
00.6
GRU
MSE
Gate
LSTM
Riman RNN
0.4
Gatel HRR
LSTT
GR17
LST3L
CateLoRD
10-3
10k
spochs
spocls
spochs
Figure 5: Shepherd results: prediction error for 100-step predictions (a) and 1-step prediction errors of the sheep's x position at the time step of reappearance (b). Success rate for capturing the sheep using MPC (c). Shaded areas show standard deviation (a-b) or standard error (c).
8

---

(b)
aucess
Simple Crosingg
Doankey--ood
BolBlushhous-R55
Lees LenCCooingS9
KeyCumiitr $311
0.5
0.6
0.5
0.5
0.5
cusilla C
wuills GaIDORDD
cuille GloCLACD
wodll GaLCOOLDD
G
20
25
25
10
25
50
25
145 foumns
165 Frames
1dd framns
146 Fromiie
5e5 tramme
Led Frattee
Figure 6: MiniGrid results: success rate in solving various tasks when GateL.ORD replaces an LSTM
(vanilla) in a PPO architecture. Shaded areas depict the standard deviation.
6.3 Long-term memorization
We hypothesized that GateLORD'ss latent state update strategy fosters the exact memorization of unobservable information, which we examine in the Shepherd task. We test the RNNs' when predicting sequences of 100 observations given the first two observations and sequence of actions. Again, we use the trained models for MPC using iCEM [51], aiming at catching the sheep by first placing cage and then pulling lever. This particularly challenging plan because the sheep's horizontal position needs to be memorized before it is occluded for quite some time (> 30 steps) in order to accurately predict and thus place the cage at the sheep's future position.
Figure 5a shows the prediction errors during training. GateLORD (A = 0.0001) continuously achieves lower prediction error than the other networks. Apparently, it is able to accurately memorize the sheep's future position while occluded. To investigate the memorization we consider the situation occurring during planning: the sequence of (past) observations is fed into the network and the prediction error of the sheep's horizontal position at the time of reappearance is evaluated (Fig. 5b). Only GateLORD reliably learns to predict where the sheep will appear when the lever activated. GRU and Elman RNNs do not noticeably improve in predicting the sheep's 's position. LSTMs take much longer to improve their predictions and do not reliably reach GateLORD's level of accuracy. This is also reflected in the success rate when the networks are used for MPC (Fig. 5c). Only GoteLORD manages to solve this challenging task with a mean success rate over 50%.
6.4 Sample efficiency in reinforcement learning
Now that we have outlined some of GateLORD's strengths in isolation, we want to analyze whether GateLORD can improve existing fL-framewsrks when it is used as a memory module for POMDPs. To do so, we consider various problems that require memory in the MiniGrid suite [48]. Previous work [42, 43, 52] used Proximal Policy Optimization (PPO) [53] to solve the MiniGrid problems. We took an existing architecture based on [52] (denoted as vanilla, detailed in Suppl. B.6) and replaced the internal LSTM module with GateLORD (A 0.01). Note, that we left the other hyperparameters unmodified.
As shown in Fig. 6 the architecture containing GateLORD achieves the same success rate or higher than the vanilla baseline in all considered tasks. Additionally, GateLORD is more sample efficient, i.c. it is able reach high success rate (Fig. 6) high reward level faster (Suppl. D.9). The difference in sample efficiency tends to be more pronounced for problems that require more training time. seems that the inductive bias of sparsely changing latent states enables GoteLORD quicker learn to encode task-relevant information, such as the pick-up of a key, within its latent states. Additional experiments in Suppl. D.10 show that this can also translates to improved zero-shot policy transfer, when the system is tested on a larger environment than it was trained on.
6.5 Explainability of the latent states
Lastly, we analyze the latent representations of GateLORD, starting with Billiard Ball. Figure 7a shows one exemplary ball trajectory in white and the prediction in red. Inputs for which at least one gate opened are outlined in black. Figure 7b shows the corresponding latent states h, relative to the initial latent state hg- GateLORD updates two dimensions of its latent states around the points of collisions to account for the changes in z- and velooicity of the ball. For X 0.01 we find on

---

aa) Billiard Ball trajectory
(b): Billiard Ball latent states
(c): Fetch Pick&Place sequence
1.5
GateLoRD
0.5
1.5
clyject
GRU
0.5
LSTM
10
20
Figure 7: Example sequences and latent states: (a) Billiard Ball trajectory for GateLORD (A = 0.01) with real positions (white). provided inputs (blue), and predicted positions (red, saturation increasing with time). The inputs for which at least one gate opened are outlined in black. (b) The latent states for the trajectory for GateLORD, GRU, and LSTM (cell states). (c) Fetch Pick&Place sequence with real (solid) and predicted (dotted) positions of gripper (top) and object (middle) and GateLORDD's latent states (bottom). Latent states are shown relative to initialization, i.e. h, hg-
average only two latent state dimensions change per sequence (see Suppl Suppl. D.I), which hints at tendency to encode z- and y-velocity using separate latent dimensions. In contrast, the exemplary latent states of the GRU and LSTM networks shown in Fig. 7h are not as easily interpretable.
For Robot Remote Control, GoteLORD (A = 0.001) updates only its latent state once it controls the robot (exemplary shown in Fig. 4d). Thus, the latent state clearly encodes control over the robot. We use the Fetch Pick&Place scenario as higher-dimenssioal problem to investigate latent state explainability when training on grasping sequences (detailed Suppl. B.5). Here, GateLORD updates the latent state typically when the object is grasped (exemplary shown in Fig. 7c). This hints at an encoding of "object transportation" using one dimension. Other RNNs do not achieve such a clear representation, neither in Robot Remote Control nor in Fetch Pick&Place (see Suppl. D.5 and D.7).
Discussion
We have introduced novel RNN architecture (GateLORD), which implements an inductive bias to develop sparsely changing latent states. The bias is realized by gating mechanism, which minimizes the La norm latent updates. In several empirical evaluations, we quantified and analyzed the performance of GateLORD on various prediction and control tasks, which naturally contain piecewise constant, unobservable states. The results support our hypothesis that networks with piecewise constant latent states can generalize better distributional shifts of the inputs, ignore spurious time dependencies, and enable precise memorization. This translates into improved performance for both model-predietin control (MPC) and reinforcement learning (RL). Moreover, demonstrated that the latent space becomes interpretable, which is important for explainability reasons.
Our approach introduces additional hyperparameter, which controls the trade-off between the task at hand and latent space constancy. When chosen in favor of explainability, can reduce the in-distribution performance while improving its generalization abilities. When the underlying system has continuously changing latent states, our regularization is counterproductive. As demonstrated by an additional experiment in Suppl. D.8, the unregularized network performs well in such cases.
Our sparsity-biasee biased gating mechanism segments sequences into chunks of constant latent activation. These segments tend to encode unobservable, hehavior-relevam states of the environment, such as if an object is currently 'under control'. Hierarchical planning and control methods require suitable, remporally-extendee encodings, such options [54, 55]. Thus, promising direction for future work is to exploit the discrete hidden dynamics of GateLORD for hierarchical, predt-pre planning.
Acknowledgments and Disclosure of Funding
The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Christian Gumbsch. Georg Martius and Martin Butz are members of the Machine Learning Cluster of Excellence, EXC number 2064/1 1-rojct number 390727645. We acknowledge
10

---

the support from the German Federal Ministry of Education and Research through the Tübingen AI Center (FKZ: 011S18039B). This research was funded by the German Research Foundation (DFG) within Priority-Progratt "The Active Self" SPP 2134-project BU 1335/11-1. The authors thank Maximilian Seitzer for the helpful feedback and Sebastian Blaes for the help in applying iCEM.
References
[1] Richard Sutton and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, Cambridge, MA, second edition edition, 2018
[2] Matthew J. Hausknecht and Peter Stone. Deep recurrent Q-learning for partially observable MDPs. arXiv preprini arXiv: 5577066577 2015. URL httpt/t////axtii http://wwwww...o...///// htrg//bs//e /aaps//100..00 1607.06527.
[3] Maximilian 1gl. Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for POMDPs. International Conference on Machine Learning, pages 2117-2126. PMLR, 2018.
[4] Pengfei Zhu, X. Li, and P. Poupart. On improving deep reinforcement learning for POMDPs. arXiv preprini arXiv: /804.06309, 2017. URL http: htap::///axxvv arxiv. arg/abs//aa 1804. 06309..
[5] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 98): 1735-1780, 1997.
[6] Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks sequence modeling. arXiv preprint arXiv: 1412.3555, barg/abee aba/ aba//4/a// 11/111111155511
2014. URL https: / http///xxxvvv////
[7] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: Foun- dations and learning algorithms. MIT press, 2017.
[8] Bernhard Schölkopf. Causality for machine learning. arXiv preprint 191 2019.
[9] Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Towards causal representation learning. In Proceedings of the IEEE, 2021.
[10] Elliot Pitis, Silviu Creager and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. In Advances in Neural Information Processing Systems 34 (NeurIPS 2020), 2020.
[11] Maximilian Seitzer, Bernhard Schölkopf, and Georg Martius. Causal influence detection for improving efficiency in reinforcement learning. In Advances in Neurul Information Processing Systems 35 (NeurIPS 2021), 2021.
[12] Dare A. Baldwin and Jessica E. Kosie. How does the mind render streaming experience as events? Topics Cognitive Science, 13(1):79-105, 2021. doi: https://doi.org/10.11111 1/tops. 12502.
[13] Martin V. Butz, Asya Achimova, David Bilkey, and Alistair Knott. Event-predictive cognition: root for conceptual human thought. Topics in Cognitive Science, 13(1):10-24, 2021. doi: https://doi.or//1...
125022.12522.
[14] Gina R. Kuperberg. Tea with milk? A hierarchical generative framework of sequential event comprehension. Topics in Cognitive Science, 13:256-298, 2021. doi: 10.1 1/tops. 1/tops.12518 12518.
[15] Gabriel A. Radvansky and Jeffrey M. Zacks. Event cognition. Oxford University Press, 2014.
[16] Jeffrey Zacks, Nicole K. Speer, Khena M. Swallow, Todd S. Braver, and Jeremy R. Reynolds. Event perception: mind-brain perspective. Psychological bulletin, 133(2):273-995 2007. doi:
[17] Martin V. Butz. Towards a unified sub-symbolic computational theory of cognition. Frontiers in Psychology, 7(925), 2016. doi: 10.3389//psy..2011600955.

---

[18] Martin V. Butz, David Bilkey, Dania Humaidan, Alistair Knott, and Sebastian Otte. Learning, planning, and control in a monolithic neural event inference architecture. Neural Networks, 117: 135-144, 2019. doi: n0.1016/).euumet221999999 201..05.00..
[19] Christian Gumbsch, Martin V. Butz, and Georg Martius. Autonomous identification and goal- directed invocation of event-predictive behavioral primitives. IEEE Transuctions on Cognitive and Developmental Systems, 13(2):298-311 June 2019. doi: 10.1 109/TCDS leee. mtg/document/8753716.6
URL https::/
[20] Dania Humaidan, Sebastian Otte, Christian Gumbsch, Charley M. Wu, and Martin Butz. Latent event- predictvei encodings through counterfactual regularization. Proceedings of the Anmual Meeting of the Cognitive Science Society, 43, 2021. URL https: org/uc/itent 5238p85g.
[21] Anna C. Schapiro, Timothy T. Rogers, Natalia I. Cordova, Nicholas Turk-Browne, and Matthew M. Botvinick. Neural representations of events arise from temporal community structure. Nat Neurosci, 16(4):486-492, April 2013. ISSN 1097-6256. URL http::/tt...... org/10. 1038/ 0038//n..4 nn. 3331.
[22] Yeon Soon Shin and Sarah DuBrow. Structuring memory through inference- event segmentation. Topics in Cognitive Science, 13:106-127, 2021. doi: 10.1111/tos..1555
[23] Martin Butz. Towards strong AI, Künstliche Intelligenz, 35:91-101, 2021. doi: 10.1007/ 0111880021-0070--0
[24] Christos Louizos, Max Welling, and Diederik P. Kingma. Learning sparse neural networks through Lo regularization. In International Conference on Learning Representations (ICLR), 2018. URL https: httpp:://peereeteeeeeee
net
[25] Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous re- laxation of discrete random variables. In International Conference on Learning Representations, ICLR'17 17, 2017.
[26] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations, ICLR 17, 2017.
[27] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3453 2013.
[28] Ronald J Williams. Simple statistical wradient-following algorithms for connectionist reinforce- ment learning. Machine learning, 8(3-4):229-2) 1992.
[29] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Leurning Representations, ICLR' 14, 2014.
[30] Nima Mohajerin and Steven L Waslander. State initialization for recurrent neural network modeling of time-series data. In 2017 International Joint Conference on Neural Networks (IJCN), pages 2330-2337. IEEE, 2017.
[31] Jimmy Ba, Volodymyr Mnih, and Koray Kavukcuoglu. Multiple object recognition with visual attention. In International Conference on Learning Representations, ICLR 15, 2015.
[32] Jürgen Schmidhuher. Learning complex, extended sequences using the principle of history compression. Neural Computation, 4(2):234-242, 1992.
[33] Jan Koutnik, Klaus Greff, Faustino Gomez, and Juergen Schmidhuber. A Clockwork RNN. In Eric P. Xing and Tony Jebara, editors. Proceedings of the 31 International Conference on Machine Learning, volume 32 of Proceedings of Machine Learning Research, pages 1863- 1871, Bejing, China, 22 24 Jun 2014. PMLR. htkkhhttp://proceeddinn mlr wll..prees/vv// koutnik14 html.
[34] Daniel Neil, Michael Pfeiffer, and Shih-Chii Liu. Phased Istm: Accelerating recurrent network training for long or event-based sequences. In Advances In Neural Information Processing Systems, pages 3882-3890, 2016.
12

---

[35] David Krueger and Roland Memisevic. Regularizing RNNs by stabilizing activations. arXiv preprint arXiv:1511.085 2015. URL https://axxix wtwpa/iwwww....a...aaaaaaaaaa aa////b///11111 /tt///111..1111 08400.
[36] Victor Campos, Brendan Jou, Xavier Giró-i Nieto, Jordi Torres, and Shih-Fu Chang. Skip mnn Learning skip state updates in recurrent neural networks. In International Conference on Learning Representations, ICLR' 18, 2018.
[37] Zhuohan Li, Di He, Fei Tian, Wei Chen, Tao Qin, Liwei Wang, and Tieyan Liu. Towards binary- valued gates for robust LSTM training. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Research, pages 2995-3004. PMLR, 10-15 Jul 2018.
[38] Thomas Hartvigsen, Cansu Sen, Xiangnan Kong, and Elke Rundensteiner. Learning selec- tively update state neurons in recurrent networks. In Proceedings of the 29th ACM International Conference on Information & Knowledge Management, pages 485-494, 2020.
[39] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401, 2014.
[40] Dzmitry Bahdanau, Kyung Hyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. In 3rd International Conference on Learning Representations, ICLR 2015, 2015.
[41] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, & ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran As- sociates, Inc., 2017. URL https: tttpse//prooeeedingg.. met.rips..c/pppe//////// (atee244347ddee11b0003lc....o..d
Paper. pdf.
[42] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In 9th International Conference on Learning Representations (ICLR 2021), May 2021. URL stttps://openre ddd?iddml.cmm1EUxyy
net/
[43] Kanika Madan, Nan Rosemary Ke, Anirudh Goyal, Bernhard Schölkopf, and Yoshua Bengio. Fast and slow learning of recurrent independent mechanisms. In International Conference on Learning Representations, ICLR'21, 2021.
[44] Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayaku- mar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International Conference on Machine Learning, pages 7487-7498. PMLR, 2020.
[45] Jeffrey L. Elman. Finding structure in time. Cogmitive Science, 14(2):179-2111 1990. ISSN 0364-0213. doi: https://doi.org// wttg://...00100/4000000000000000 sci sciencedirect. com/ saaasccnnceeeeeeeeeeeeeeeeeeeece htt.a////aaatttee///////////0000
URL http:: //:w://wwww
[46] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, ICLR' 14, 2014.
[47] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. OpenAl gym. arXiv preprint arXiv:1606.01540 2016.
[48] Maxime Chevalier-Boisvert Lucas Willems, and Suman Pal. Minimalistic gridworld environ- ment for openai gym. httpp::/ggithhe www.cooo//////e httm///xmmmec///mmmmmm ggggggmmmmiiiii...mm 2018.
[49] Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neurul Information Processing Systems, vol- ume 28. Curran Associates, Inc., 2015. URL https:/ htt.tttppp/ppppeepiegeeeeeeeeeee neurips. cwtuurps..ccpappee.. 2015
pdf.
13

---

[50] Anirudh Goyal Lamb, Alex M. Ying Zhang, Saizheng Zhang, Aaron C Courville, and Yoshua Bengio. Professor forcing: A new algorithm for training recurrent networks. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016. URL / neurips. wwtpppppe///1///1eeeeeee 00T0///2200000000000000000000000
[51] Cristina Pinneri, Shambhuraj Sawant, Sebastian Blaes, Jan Achterhold, Joerg Stueckler, Michal Rolmek, and Georg Martius. Sample-efficiett cross-entropy method for real-time planning. In Conference Robot Learning 2020, 2020. URL https::/corlcooooo github. io io/paper_ 217.
[52] Maxime Chevalier-Boisvert Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Sa- haria, Thien Huu Nguyen, and Yoshua Bengio. BabyAI: A platform to study the sample efficiency of grounded language learning. In International Conference on Learning Representa- tions, ICLR' 18, 2018.
[53] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint (rxXiv:707.0000 2017.
[54] Richard S. Sutton, Doina Precup, and Satinder Singh. Between MDPs and semi-MDPs: A framework for temporal abstraction in reinforcement learning. Artificial Intelligence, 1 112(1)):
181-211, 1999. ISSN 0004-3702. doi: https::/doo.oor/I00 000.1000066000004070000900000000 URL https: http::///ww..siieedddreettt www.secciiieiccee...eo caltttt.comcccii wwwmm//cince//eeeeeeeeeeeii wttpa////arraae/////////00000000
[55] Andrew G Barto and Sridhar Mahadevan. Recent advances in hierarchical reinforcement
learning. Discrete event dynamic systems, 13(1):41-77, 2003.
[56] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty training recurrent neural networks. In Proceedings of the 30th International Conference on International Con- ference on Machine Learning Volume 28, ICML'13, 13, page III-1310-1 1318. JMLR.org. 2013.
[57] Cristina Pinneri, Shambhuraj Sawant, Sebastian Blaes, and Georg Martius. Extracting strong policies for robotics tasks from zero-order trajectory optimizers. In International Conference on Learning Representations, ICLR 21, 2021.
[58] Lu Lu. Dying ReLU and initialization: Theory and numerical examples. Communications in Computational Physics, 2865) Jun 2020. ISSN 1991-7120. doi: 10.4208/cicp. oa-2020-0165. URL http: //t::////x doi..orgooo.o...o............... 100g/ 4208/22c0ppp OA 2020-0165.
14


Reviews:
[{"review_id": "t7qktGFCD2I", "reviewer": "Reviewer_yJcq", "summary": "This work proposes an RNN with a gating mechanism that encourages its latent state to have sparse updates by incorporating stochastic gating units trained using an approximation of L0 regularization of the state updates. The authors propose to use this RNN in partially observable problems in which the latent dynamics are known to change slowly and long-term memory is important. The authors show in a number of partially observed tasks that their proposed RNN and training objective outpeforms well known RNN architectures when predicting the environment dynamics.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "* Originality and significance\n\nSlow changing dynamics is a pervasive property of many time series and (PO)MDPs.  While designing architectures and training algorithms that reflects this property is not novel, it still remains an open research problem. The method proposed, an RNN that closely resembles LSTMs, is very closely related to some existing techniques as stated in the related work, with the main difference being in how the new gating units are modelled and trained. The authors note that these prior works have not been applied to POMDP problems but instead to conventional classification and language-based tasks. However, because the authors do not compare with these works, it's unclear how those techniques (see references [27, 28, 29]) would fair in POMDPs problems. The authors claim that a GRU does not include an output that is not its state (making it difficult to predict changing observations when the state remains constant), however that can be simply remedied by using the observation in addition to the state in the f_post component.\n\nEvaluating these ideas in POMDPs is original, and the experiment section show a clear advantage of the RNNs sparsely changing bias over the baselines which is definitely interesting. Since the main novelty is the application to POMDP problems (and not as much the method itself), the selection of environments should be more extensive. The environments' partial observability and dynamics are fairly simple (e.g. they do not incorporate stochastic transitions). This paper would thus be strengthed if the authors compare relevant work in modelling the dynamics in challenging POMDPs and discussed the method's advantages and limitations (see Igl et al. 2018, Hafner et al. 2019 for variational methods for example).  If the main contribution of this work is to use the proposed RNN for long-term memory of unchanged environment states, then perhaps the paper would benefit from presenting it as a general module that can replace the core RNN in more complex architectures. For instance, many variational methods use an LSTM as the memory back-bone. Would their performance improve with this?\n\n* Quality\n\nThe proposed RNN is tested in terms of next step prediction as well as its usefulness for planning. The second aspect is important to assess the quality of the predictions when unrolled over time. \nThe experiments succesfully show the benefits of sparse updates in different environments and settings (e.g. with generalization to unseen trajectories), as well as show the explainability of the learned states. There's also a good discussion on the trade-offs of predictive performance and explainability (sparseness) which is important for the reader to know. Results also seem robust as the error bars are relatively small.\n\nThe L0 regularization method is well derived and closely related to [16]. However this work, contrary to the relaxation of the intractable L0 regularization, does not minimize the probability of non-zero activations as in [16] (which is differentiable) as a proxy of L0, but instead directly incorporates a non-differentiable term based on the sum of gate values, which is approximated with the straight-through estimator. It's unclear why the authors choose this particular solution which is more biased (due to using the straight-through estimators in *both* terms of the loss), and also what practical consequences this has. \nI would also like to see an ablation study where a simpler version of this idea is tried: a sigmoid gate is used with L1 regularization. In other words, to what extent is it necessary to use stochastic gates?\n\n* Clarity\n\nThe paper is well-written and the methods are generally clear, well explained and motivated.  There are some parts where certain details are explained in the appendix but not referenced from the paper so it seemed at first a bit cryptic. For instance, in line 93 they refer to a \"variation of the straight-through estimator\" without specifying which variation. Similarly, many interesting ablation studies are in the appendix without any reference to them.\n\nSome other minor issues:\n- typo in line 75 (combine > combined).\n- line 248, I believe the authors mean that a higher value results in *more* latent state changes?\n\nReferences:\n\nMaximilian Igl, et al., Deep Variational Reinforcement Learning for POMDPs. ICML 2018\n\nDanijar Hafner et al., Learning Latent Dynamics for Planning from Pixels. ICML 2019", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have been fairly clear about the limitations and trade-offs of their method; and as a paper focused on methodology, it does not seem to have any direct social and ethical implication.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "TzFbAyYNo-c", "reviewer": "Reviewer_qktj", "summary": "The manuscript presents a novel recurrent neural network architecture with an inductive bias for learning sparsely changing hidden state representations. This seems to be achieved by the proposed gating scheme using the novel rectified tanh activation function for the update gate values. This activation function yields true zero activations producing an exact copy of the corresponding value from the previous hidden state. By applying L0 regularization to the state-updates, the network is incentivized to minimize the number of modified dimensions in the state.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The presentation is very clear and detailed. The experiments highlight several contributions: 1) GateL0rd develops stable, sparsely updated latent state representations. 2) making the output also depend directly on the input, frees the hidden state to focus on information inferred from past time steps. To summarize the points made in the sections below, I think the work proposes an intriguing alternative to commonly used RNN architectures, but the experimental evaluation has some weaknesses unless I’m misinterpreting some observations.\n\nStrengths\n1. The proposed regularized ReTanh gating and direct dependence of the network output on the input seem like very appropriate solutions to weaknesses of other (gated) RNN architectures. Especially the potential disentanglement of observable and unobservable information might help in reinforcement learning with an unknown reward function.\n1. The sparse state changes should be very useful for event-based sequence modeling.\n1. GateL0rd seems to outperform other commonly used RNN architectures.\n1. The manuscript is very well written. A lot of questions that came up while reading were answered either a few paragraphs later or in the supplementary material. For instance, I was a bit sceptical at first about the choice of the ReLU gradient estimator, but Section C.2 cleared that up for me.\n\nLimitations/Questions\n1. How do you think GateL0rd would perform if applied to environments with very high-dimensional visual observations, such as MuJoCo, Atari or ProcGen (of course using an appropriate encoder network for the frames)? I think that GateL0rd certainly has a better prior for disentangling factors of variation in the environment than GRUs or LSTMs. However, I also think that it may be difficult to learn this disentanglement, if at the same time a compatible input representation has to be discovered given high-dimensional observations.\n1. Footnote 3: I think that tanh(s/2) = 2 * sigmoid(s) - 1. \n1. Can you quantify the impact of $f_{init}(x)$ compared to setting $h_0=0$?\n1. The testing error curves in Figures 4 and 5 look like training did not  yet converge. If that is the case, the networks should be trained longer since the focus of the manuscript does not seem to be on comparing the sample efficiency of different architectures.\nIf training did converge, I’d suggest presenting the loss curves as well.\n1. Can the test prediction error going up for GRU and LSTM during training really be attributed to the networks’ overfitting to observable information and spurious correlations? The variance of these curves seems to be very large which is often a symptom of optimization issues (e.g. exploding gradients for some runs)! \n1. In the shepherd task, suppose the sheep was allowed to move horizontally while occluded with horizontal motions sampled from some symmetric zero-mean distribution. If the motions were part of the observation, do you think GateL0rd would be able to learn to modify the memorized initial position accordingly in its hidden state? If the motion information were not observed, how would GateL0rd learn to account for the uncertainty?\n1. Time indices are not consistent. In line 106 the mapping is (y_t, h_{t+1}) = f(x_t, h_t), while in line 71 it is (y_t, h_t) = f(x_t, h_{t-1}). Also in Equation 2, I believe the regularization should be applied to delta h_t.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See main review.", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "JeknaC2uzlH", "reviewer": "Reviewer_EaBh", "summary": "This paper proposes a new recurrent neural network architecture (RNN) called GateL0RD, designed to learn sparsely changing latent variables of sequential data. The authors hypothesize that latent variables should be updated sparingly only if significant changes happen in the considered environment. To enforce this hypothesis in RNN, GateL0RD uses a new gating function of Rectified Tanh to output exactly zero values. For actual implementation, the considered prediction task is regularized to minimize the number of positive elements in the sampled gate inputs. Ablation study shows that GateL0RD improved prediction accuracy and planning performance compared with LSTM or GRU in several POMDP environments.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "(1) Originality \\\nThis work builds a new activation function (ReTanh) on a gating function in usual LSTM/GRU architectures where sigmoid and $\\tanh$ are widely used. \n\n(2) Quality \n1) Simple but practical implementation of $L_0$ regularization makes this work scalable.\n2) Various experimental results and the corresponding ablation study are provided, including the interpretability.\n3) Main concern of this work is the support of the hypothesis that latent variables should change sparsely. It looks like there is no theoretical support or guarantee on the optimality of this approach, though several examples in the real-world and human recognition area are mentioned.\n\nQ1. Is there any more substantial explanation or related work supporting the effectiveness of the hypothesis? \\\nQ2. Does the proposed method perform better than $L_1$, $L_2$ regularization methods?\n\n(3) Clarity \\\nThis paper is well written and easy to read. The strong points of this work are well organized regarding items: accuracy, generalization, long-term memorization, and interpretability.\n\n(4) Significance \\\nAs mention in the Discussion section, the proposed regularization may not be effective when the considered partial domain has severely changing characteristics. \n\nQ3. Is there any backup experiment (any specific environment) to check this aspect? Can this regularization be used in general in partially observable domains? Setting $\\lambda=0$ and asserting that \"the unregularized network performs well.\" is not persuasive. The proposed architecture without the regularization is not fundamentally different from LSTM / GRU. Also, note that we usually do not know a priori how much the considered environment changes.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the authors have addressed the limitations and societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "12"}, "extra_text": {}}, {"review_id": "EdiFFGCrwO4", "reviewer": "Reviewer_sDuQ", "summary": "This paper presents a sequence prediction method that incorporates a particular inductive bias -- infrequency of latent state updates -- as measured by an L_0 term on latent state differences. This term is approximated and made more easily optimizable using a method similar to [16], which examined L_0 regularization in non-sequential domains. The paper conducts experiments on several low-dimensional POMDP tasks and then draws the following conclusions: the method generalizes better to out-of-distribution inputs in several partially observable domains than other commonly used RNNs, exhibits long term memory in a domain, and exhibits explainability (in its latent states), sometimes at the cost of reduced prediction performance.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Opinion summary**: The paper offers a well-motivated method for sequence prediction in partially-observed domains that contain dynamics with significant changes. The proposed method is straightforward and clear. However, the method is of questionable originality (see below). The experiments are conducted relative to basic common NN-based sequence-prediction methods, which do show significant improvements, but lack comparison to previously-proposed methods for incorporating the same inductive bias (infrequency in latent state updating). Thus, we cannot conclude that this method is better than a relevant set of previous methods for accomplishing the same goals. There are a few other issues with the experiments that could be improved.\n\n**The paper has questionable originality, as prior work accomplishes a similar goal in a different way**: there are missing comparisons to existing methods -- [25,26] use predefined time scales for gate opening, which may be sufficient for high performance on the considered tasks, and [27,28,29] all use binarized update gates, and are thus quite related to the proposed method. The proposed method is not clearly different enough from these methods, and is therefore of questionable novelty. I think empirical results that compare the methods are needed -- without them, we cannot understand situations in which the proposed method is better or worse than prior work that attempts to accomplish the same goal.\n\n**The interpretability is not quantitatively supported:** The method pitches sparsity in gate opening as both a way to learn a more performant system, as well as a way to understand when the dynamics of the environment are significantly changing. On the latter point -- Figs 4d and 6b illustrate examples where the gate opening makes qualitative sense (the paper describes these as \"interpretable\"). These results provide intuition, but are not evidence of how reliable the method is at representing \"significant and systematic\" changes with gate openings. It would be interesting to have a better understanding of how reliable it is, for instance by (automatically?) labeling each point in a sequence as either a \"significant and systematic\" change or not (like the black-outlined positions in Fig 6a and the moving-robot positions in Fig 4d), and then reporting binary classification results (e.g. precision, recall, accuracy) using the classifier p(systematic change=1 | any gate open=1)=1; p(systematic change=0 | any gate open=0)=1.\n\n**Other issues with experiments (ordered by decreasing priority):**\n- L270 \"only encodes unobservable information within its latent state\" this claim has, at best, anecdotal evidence (Fig 4d). It is a significant claim to state generally, so either temper it or provide sufficient evidence to support it.\n- The $\\lambda=0$ results in Fig 3d requires further explanation. Why does the version with no regularization exhibit the fewest latent state changes (L248 says \"a higher value results in fewer latent state changes.\")?\n- Fig 3c is too zoomed-out. The differences between the nonzero lambda experiments are difficult to see.\n- Fig 15 is referenced in the text but appears in the supplement.\n\n---\nPost-response update: My above concerns have been addressed.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations were appropriately described; the societal impacts were described inside the checklist (which are essentially a statement that there are no ethical concerns specific to the work). I agree with that sentiment, but it should be incorporated either into the main text or into the appendix, rather than the checklist, which will not be present in a potential final version(?)", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a recurrent architecture with an inductive bias (via gating and L0 regularization) that encourages sparse updates to its state. This is hypothesized to lead to improved generalization, and is empirically evaluated in a number of partially observable environments.

The reviewers were initially sceptical, but the discussion with the authors has allayed most concerns, so now the consensus is towards acceptance. In general, the reviewers agree that the paper is well written and the contribution is interesting (albeit somewhat incremental). Overall, I'm happy to recommend acceptance.
Decision: Accept (Poster)