Episodic Memory for Learning Subjective-Timescale Models
Authors: Alexey Zakharov
Matthew Crosby
Zafeirios Fountas
Venue:  (2021)
Keywords: Episodic Memory
Time Perception
Active Inference
Model-based Reinforcement Learning

Abstract: In model-based learning, an agent’s model is commonly defined over transitions between consecutive states of an environment even though planning often requires reasoning over multi-step timescales, with intermediate states either unnecessary, or worse, accumulating prediction error. In contrast, intelligent behaviour in biological organisms is characterised by the ability to plan over varying temporal scales depending on the context. Inspired by the recent works on human time perception, we devise a novel approach to learning a transition dynamics model, based on the sequences of episodic memories that define the agent's subjective timescale – over which it learns world dynamics and over which future planning is performed. We implement this in the framework of active inference and demonstrate that the resulting subjective-timescale model (STM) can systematically vary the temporal extent of its predictions while preserving the same computational efficiency. Additionally, we show that STM predictions are more likely to introduce future salient events (for example new objects coming into view), incentivising exploration of new areas of the environment. As a result, STM produces more informative action-conditioned roll-outs that assist the agent in making better decisions. We validate significant improvement in our STM agent's performance in the Animal-AI environment against a baseline system, trained using the environment's objective-timescale dynamics.

Under review as a conference paper at ICLR 2021
EPISODIC MEMORY FOR LEARNING SUBJECTIVE- TIMESCALE MODELS
Anonymous authors Paper under double-blind review
ABSTRACT
In model-based learning, an agent's model is commonly defined transitions between consecutive states of an environment even though planning often requires reasoning over multi-step timescales, with intermediate states either unnecessary, or worse, accumulating prediction error. In contrast, intelligent behaviour in bi- ological organisms is characterised by the ability to plan over varying temporal scales depending on the context. Inspired by the recent works on human time perception, we devise a novel approach learning a transition dynamics model, based on the sequences of episodic memories that define the agent's 's subjective timescale - over which it learns world dynamics and over which future planning is performed. We implement this in the framework of active inference and demon- strate that the resulting soojective-timesaal model (STM) can systematically vary the temporal extent of its predictions while preserving the same computational ef- ficiency. Additionally, we show that STM predictions are more likely to introduce future salient events (for example new objects coming into view), incentivising exploration of new areas of the environment. As a result, STM produces more informative action-condittoned roll-outs that assist the agent in making better de- cisions. We validate significant improvement in our STM agent's performance in the Animal-AI environment against baseline system, trained using the environ- ment's 's tijestlve-timescale dynamics.
INTRODUCTION
An agent endowed with model of its environment has the ability to predict the consequences of its actions and perform planning into the future before deciding on its next move. Models can allow agents to simulate the possible action-conditioned futures from their current state, even if the state was never visited during learning. As result, model-based approaches can provide agents with better generalization abilities across both states and tasks in an environment, compared to their model-free counterparts (Racanière et al. 2017: Mishra et al. 2017).
The most popular framework for developing agents with internal models model-based reinforce- ment learning (RL). Model-based RL has
great progress in recent years, with number of proposed architectures attempting to improve both the quality and the usage these models (Kaiser al. 2020: Racanière al.) 2017 Kansky al. 2017 Hamrick 2019). Nevertheless, learning internal models affords a number of unsolved problem mroblems. The central one of them is model-bias, in which the imperfections of the learned model result in unwanted over-optimism and sequential error accumulation for long-term predictions (Deisenroth & Rasmussen 201I). Long-term predictions are additionally computationally expensive in environments with slow temporal dynamics, given that all intermediary states must be predicted. Moreover, slow world dynamics"] can inhibit the learning of dependencies between temporally-distar distant events, which can be crucial for environments with sparse rewards. Finally, the temporal extent of future predictions is limited to the objective timescale of the environment over which the transition dynamics has been learned. This leaves little room for flexible and context-dependent planning over varying timescales which is characteristic to animals and humans (Clayton et al. 2003 Cheke & Clayton 2011 Buhusi & Meck 2005).
The final issue exemplifies the disadvantage of the classical view on internal models, in which they are considered to capture the ground-truth transition dynamics of the environment. Furthermore,
Worlds with small change in state given an action

---

Under review as a conference paper at ICLR 2021
in more complex environments with first-person observations, this perspective does not take into account the apparent subjectivity of first-person experiences. In particular, the agent's 's learned repre- sentations of the environment's ' transition dynamics implicitly include information about time. Little work has been done to address the concept of time perception in model-based agents (Deverett et al.) 2019). Empirical evidence from the studies of human and animal cognition suggests that intelligent biological organisms do not perceive time precisely and do not possess an explicit clock mechanism responsible for keeping track of time (Roseboom et al. 2019; Sherman al. 2020 Hills 2003). For instance, humans tend to perceive time slower in environments rich in perceptual content (e.g. busy city), and faster in environments with little perceptual change (e.g. empty field). The mechanisms of subjective time perception still remain unknown; however, recent computational models based on episodic memory were able to closely model the deviations of human time perception from veridical perception (Fountas et al. 2020b).
Inspired by this account, this work we propose subjective- model (STM), an alterna- tive approach to learning transition dynamics model, by replacing the objective timescale with subjective one. The latter represents the timescale by which an agent perceives events in an envi- ronment, predicts future states, and which is defined by the sequences of episodic memories. These memories are accumulated on the basis of saliency (i.e. how poorly an event was predicted by the agent's transition model), which attempts to mimic the way humans perceive time, and resulting in the agent's ability to plan over varying timescales and construct novel future scenarios.
We employ active inference as the agent's underlying cognitive framework. Active inference is an emerging framework within computational neuroscience, which attempts to unify perception and action under the single objective of minimising the free-energy functional. Similar to model-based RL, an active inference agent relies almost entirely on the characteristics and the quality of its inter- nal model to make decisions. Thus, it naturally susceptible to the previously mentioned problems associated with imperfect. objective- abjective-timescale models. The selection of active inference for the pur- poses of this paper is motivated by its biological plausibility as a normative framework for under- standing intelligent behaviour (Friston et al., 2017a 2006), which is in line with the general theme of this work. Furthermore, being rooted in variational inference, the free energy objective generates distinct separation between the information-lle -hooortion-thoorttiic quantities that correspond to the different components of the agent's model, which is crucial to define the memory formation criterion.
We demonstrate that the resulting characteristics of STM allow the agent to automatically perform both short- and long-term planning using the same computational resources and without any explicit mechanism for adjusting the temporal extent of its predictions. Furthermore, for long-term predic- tions STM systematically performs temporal jumps (skipping intermediary steps), thus providing more informative future predictions and reducing the detrimental effects of one-step prediction error accumulation. Lastly, being trained on salient events, STM much more frequently imagines futures that contain pristemically-surrritttti events, which incentivises exploratory behaviour.
2 RELATED WORK
Model-based RL. Internal models are extensively studied in the field of model-based RL. Using linear models to explicitly model transition dynamics has achieved impressive results in robotics
(Levine& & Abbeel 2014a] Watter al. 2015 Bagnell & Schneider) [001]] Abbeel al. 2006
Levine & Abbeel 2014b Levine et a 2016 Kumar et 2016). In general, however, their applica- tion is limited low- low-dimensional domains and relatively simple environment dynamics. Similarly, Gaussian Processes (GPs) have been used (Deisenroth & Rasmussen 2011] Ko et al.) (007). Their probabilistic nature allows for state uncertainty estimation, which can be incorporated in the plan- ning module to make more cautious predictions; however, GPs struggle to scale to high-dimensee data. An alternative and recently more prevalent method for parametrising transition models is to use neural networks. These are particularly attractive due to their recent proven success in a variety of domains, including deep model-free RL (Silver et al. (0017), ability to deal with high-dimensional data, and existence of methods for uncertainty quantification (Blundell et al. 2015 Gal & Ghahra- mani 2016). Different deep learning architectures have been utilised including fully connecco neu- ral networks (Ngabandd et al. 20188] Feinberg et al., 2018 Kurutach et al. 2018) and autoregressive
models (Ha & Schmidhuber] 2018 Racanière et al. 2017; Ke et al. 2019), showing promising re-
sults in environments with relatively high-dimensional state spaces. In particular, autoregressive
2

---

Under review as a conference paper at ICLR 2021
architectures, such as Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber 1997, capable of modelling non-Markovian environments and learning temporal dependencies. Never- theless, LSTMs still limited in their ability to learn relations between temporally-distar events, which is exacerbated in environments where little change occurs given an action.
Uncertainty quantification using ensemble methods (Kalweit & Boedecker 2017 Clavera et al., 2020 Buckman et al. 2018) or Bayesian neural networks (McAllister & Rasmussen 2 Depeweg al., 2017) have been proposed tackle model bias and sequential error accumulation. Other works have focused on techniques to create more accurate long-term predictions. MMshra et al. (2017) used a segment-basee approach predict entire trajectories once in an attempt to avoid one-step prediction error accumulation. A work by Ke et al.] (2019) used an autoregressive model and introduced a regularising auxiliary cost with respect to the encodings of future observations, thus forcing the latent states to carry useful information for long- long-horizon predictions. In contrast, the work presented in this paper re- re-fouuses the objective from attempting to create better parametrisation techniques mitigating methods to simply transforming the timescale over which the dynamics of an environment is learned. As will be seen, our approach can lead to more accurate and efficient long-tere predictions without compromising agent's ability to plan over short -orizoos.
Episodic Memory. In neuroscience, episodic memory is used to describe autobiographical memo- ries that link collection of first-person sensory experiences at a specific time and place (Tulving) 1972). Past studies in the field suggest that episodic memory plays an important role in human learn- ing (Mahr & Csibra 2017). and may capture wide range of potential functional purposes, such as construction of novel future scenarios (Schacter et al. 2007 2012 Hassabis et al. (007), mental time-travel (Michaelian 2016) or assisting in the formation new semantic memories (Greenberg & Verfaellie 2010). A recent computational model of episodic memory (Fountas et al. relates the human ability to estimate time durations.
2020b) also
The application of episodic memory in reinforcement learning has been somewhat limited. Some works have employed simple forms of memory improve the performance of deep model-free RL agent via experience replay (Mnih et al. 2015) Espeholt et al. 2018 [Schaul et al.) (016)) However, these methods do not incorporate information about associative or temporal dependencies between the memories (Hansen al., 2018). Read-write memory banks have also been implemented along-
side gradient-basee systems (memory-augmented neural networks) for assisting learning and pre-
diction (Graves et al. 2014 2016 Hung et al. 2019 Oh et al. 2016 Jung et al.) 2018 Further, episodic memory has been used for non non-nnnn-naaaae parametric O-fucction approximation (Blundell et al. | 2016)] Pritzel et al.) 2017 Hansen et al. 2018 Zhu et al. 2020). It has also been proposed to be used di- rectly for control as a faster and more efficient alternative to model-based and model-free approaches in RL, such as instance-based control (Lengyel & Dayan 2007 Botvinick et al.) 2019 Gershman & Daw, (017)) and one-shot learning (Kaiser et 2017]. In contrast, our paper considers a novel way of using episodic memories in defining the agent's subjective timescale of the environment and training transition dynamics model over the sequences of these memories.
Active Inference. Until now, most of the work on active inference has been done in din-enmeniional and discrete state spaces (Friston et al. [2015]2017b]] c d). Recently, however, there has been rising interest in scaling active inference and applying to environments with continuous and/or large state spaces (Fountas et al., 2020a] Tsscanntz et al. 2019 Çatal et al., 2019: Millidge| (209:) Ueltzhöffer] 2018). Although these works used deep learning techniques, their generative models have so far been designed to be Markovian and trained over the objective timescale of the environment.
3 BASELINE ARCHITECTURE
We take the deep active inference system devised by Fountas et al. (2020a) as the starting point with a few architectural and operational modifications. The generative model this baseline agent is defined as p(O1.,, St.r, A1.r; 0), where denotes latent states at time t, o, (visual) observations, a actions, and 9 = {9.. 02) the parameters of the model. Se is assumed to be caausian-distribiin with diagonal covariance, o follows Bernoulli and a1; categorical distributions. For single time step, as illustrated in Figure ITA, this generative model includes two factors, a transition model p(Se|St-1 1, a 1;; 0.) and latent state decoder p(or|se; 0.) parametrised by feed-forwall neural net-
3

---

Under review as a conference paper at ICLR 2021
works with parameters 8, and f.r respectively. We modify the transition model from the original study to predict the change in state, rather than the full state
The agent also possesses two inference networks, which are trained using amortized inference: ha- bitual network ( 0a) and observation encoder (8,; Ø,) parametrised by du and Φ., respectively. The habitual network acts as model-free component of the system, learning to map inferred states directly to actions. Following Fountas et al. (2020a), the variational free energy for an arbitrary time-step & is defined as:
Fe = Egts)) [log p(oz (s;; 8o)l
(Ia)
+ DKL (u(s;; 02) - a 0s)l
(Ib)
+ gg(s)) [DEL (g(a;: (.)]]p(a)]]]
(Ic)
where p(a) = Ex.,, Res. ((T)) is the summed probability of all policies beginning with action a. All the divergence terms are computable in flosed-form given the assumption about Gaussian- and Bernoulli Bernoulli-dsstrinee distributed variables. Finally, the expected free energy (EFE) of the generative model up to some time horizon can be defined as:
T
G(T) = Σ G(T, = E Eg q(, 0//=
0|)))
(2)
where q q(,,, 8 0 T) and p(or, 87, m)
(p(())p(()))))))) 0r; T.)
To make expression 2 computationaly feasible, is decomposed such that,
G(ā,T) = (g(0)))g(s- 0,x)g(o- 0.,, [lo//((,|TT) [+] + (g(())) [Egt,, (0,m) H (s H [x)
(3)
+ (gg 0,x) (o, 0,x))
where expectations can be taken by performing sequential sampling of 0, 8, and o, and entropies are calculated in closed-form using standard formulas for Bernoulli and Gaussian distributions. Network parameters, 0, are sampled using Monte Carlo (MC) dropout (Gal & Ghahramani] (0166).
The system also makes use of top-down attention mechanism by introducing variable w, which modulates uncertainty about hidden states, promoting latent state disentanglement and more ef- ficient learning. Specifically, the latent state distribution is defined Gaussian such that 8 ~ (s; μr, E/w), where ls and E are the mean and the diagonal covariance, and w a decreasing logistic function over the divergence DRL. (g(a; (p)))p(aa))
Finally, action selection aided with Monte Carlo tree search (MCTS), ensuring a more efficient trajectory search. Specifically, MCTS generates weighted tree that used sample policies from the current timestep, where the weights refer to the agent's 's estimation of the EFE given a state-action pair, G(s, a). The nodes of the tree are predicted via the transition model, p(Se St - 0.). At the end of the search, MCTS is used construct the action prior, p(a) N 8) Σ3 ((a s),
where (s,, a) is the number of times action a has been taken from state
The baseline agent is trained with prioritised experience replay (PER) (Schaul et al. 2016) to mit- igate the detrimental consequences of on-line learning (which was used the original paper), and to encourage better object-centric representations. The details of the baseline implementation and training with PER can be found in Appendices B. and B.2 respectively.
SUBJECTIVE-TIMESCA MODEL
We introduce subjective-timescae model (STM) that records sequences of episodic memories over which a new transition model is trained. As such, the system consists of memory accumulation system to selectively record salient events, a simple action heuristic to summarise sequences of actions between memories, and an autoregressive transition model.
This has largely become common practice in the field of model-based RI. Nagahandi al. 2018), im- proving algorithm efficiency and accuracy especially environments with slow temporal dynamics.
4

---

Under review as a conference paper at ICLR 2021
P(a)
A
B
P(a)
41-1
ar-1
9-918111,4-1;;,
aa-11 va)
h
hr
p(s, 1 Ohr)
8p-1
S1 a(s; +b)
Sr-1
p(o, s 0.)
g(s; Φ,)
ps: 0z)
O
Or
Figure 1: A. Baseline generative model. hidden states h introduced by an LSTM.
B. STM generative model with additional deterministic
We define a ground-truth sequence as a sequence of all states experienced in an environment during single episode, Sg {so, 81, $2, --., ST and S-sequence (subjective sequence) sequence of states selectively picked by our system, and over which the new transition model would be learned, J.y Each unit in an S-sequence called an episodic memory and consists of set of sufficient statistics, s = {Ms, T.). where μ4, and o, are mean and variance vectors of mqussian-distribus
SrS: S. = {STL 5 ...
state s, respectively. Additionally, each episodic memory contains a reference to its preceding (parent) episodic memory and all actions until the next one. The process of recording S-sequences is called memory accumulation.
4.1 MEMORY ACCUMULATION
Previous work time perception and episodic memory (Fountas al., 2020b) employed saliency of an event, or the generative model's prediction error, as the memory formation criterion. Selection of this criterion is informed by the experimental evidence from neuroscience episodic memory (Greve et al., 2017 Jang et al. 2018 Rouhani et al., 2018). Inspired by this account, our memory accumulation system employs the free energy of the objective-timescale transition mode (Eq Ib) as measure of event saliency, and forms memories when pre-defined threshold is exceeded.
To train STM, an active inference agent moves in the environment under a pre-trained generative model described in Section During this process, each transition is evaluated based on the objective transition model free energy. DKL [g(s;: ( 1; 1; 0.)], which represents the degree of surprise experienced by the transition model upon taking an action. If the value of the free energy exceeds a pre-defined threshold, e, a memory is formed and placed into an S-sequence. the end of each episode, the recorded S-sequence is saved for later use.
We can categorise the transitions that cause higher values of transition model free energies into two main groups: epistemic surprise and iedereiimerrioction surprise. The former refers to transitions that the model could not have predicted accurately due to the lack of information about the current state of the environment (e.g. objects coming into view). The latter refers the main bulk of these high prediction-error transitions and stems from the inherent imperfections of the learned dynamics. Specifically, less frequently-occurrin would systematically result in higher compounded transition model errors, given that these would be characteristic of more complex scenes. As will become apparent, the presence of these two categories in the recorded S-sequences results in the model's 's ability to vary its prediction timescale based on the perceptual context and systematically imagine future salient events.
observations with richer combinatorial structure
A transition dynamics model is necessarily trained with respect to actions that an agent took to reach subsequent states. However, STM records memories over an arbitrary number of steps, thus leaving action sequences of variable length. For the purposes of this paper, we implement simple heuristic to summarise agent's trajectories, which enough to provide STM with the necessary information
CComponents of the total free energy correspond to measure of belief update for each of the networks, and therefore, loosely speaking, quantify the prediction error generated by each of the respective system con- stituents: autoencoder (Eqs Ta lbt, objective-timesaal transition model (Eq Tb), and habitual network (Eq Ic
5

---

Under review as a conference paper at ICLR 2021
fast
subjective timeccaee
slow
Buffer
Ground-truth sequence
Training STM Transition Model p(s ar, h:
S-sequence
8r. (1
1..
on-line fashion, pick states that exceeded the transition model free energy threshold
Save the S-sequence in the Buffer
Sample batches of S-sequences for training STM transition model
Figure 2: STM pipeline. (A) As the agent moves through the environment, states a that exceeded a pre-defined threshold are recorded along with all successive actions in S-sequence. (B) S- sequences are saved in a buffer at the end of each episode. (C) S-sequences are sampled for training a subjective-timesale transition model.
learn action-conditioned predictions. We do it by estimating the angle between the agent's initial position and its final position at the time-step of the subsequent memory. Full details of this heuristic can be found in Appendix B.4
4.2 TRANSITION DYNAMICS MODEL.
As mentioned, S-sequences characterised by the presence of pesstemmccally-suprssiiiii and salient events squeezed together in the recorded episodes. As a result, training on these sequences is more conducive for learning temporal dependencies between important states. For this reason, train an LSTM model over the S-sequences, which utilises internal memory states to store information about preceding inputs. In architecture, LSTM calculates hidden state h, subjective time r using a deterninistic mapping.
h = for (ST, 0r, Cr, 11- (x, + h--1Uh bk),
where 8, and u, are the latent state and action taken subjective time respectively, v, is the con- catenated vector of s, and a,, and 6% = {W,, Un, by are deterministic LSTM model parameters. Importantly, function fe. is deterministic and serves only encode information about preceding steps into the hidden state the LSTM. This hidden state h, is then mapped to a latent state 8 at via feed-forward neural network with random-variable parameters,
the next subjective time On.. using p(s,/ -1; (na) with MC dropout. The parameters of both of the networks are trained via backpropagation with a loss function defined
T C C= (
an, 6no)
(5)
The new generative model of observations is shown in Figure Because the mapping of LSTM deterministic, the formulation the variational free energy remains intact with the exception of the second term that now includes the state prediction produced by the network 1/, 0n.) conditioned on the hidden state of the LSTM,
F, = Eg(er) [log (lors(or)sr:))) (3r
+ DEL (g(s+; & ( 0he)]
(6)
+ By(er) [Dx [g(a:; (b)())()))]]]
Architectural and training details of the model can be found in Appendix B.3 The source code will be made available after the review process.
5 EXPERIMENTS
The Animal-AI (AAI) environment a virtual testbed that provides an open-ended sandbox training environment for interacting with a 3D environment from first-person observations (Crosby et al.

---

Under review as a conference paper at ICLR 2021
1200
160
SIM-MCIS
Cumulative Reward
Numbar of Rewards
5h 57min
Baseline-MCTI
800
20min
Baseline-MPC
600
400
200
4h 32min
Timestep (x10%)
10
Green
Yellow
Red
Figure 3: Experimental Results. (A) Cumulative rewards collected by the agents. can be seen that the STM-MCTS agent shows improved performance even when compared with the computationally: expensive Baseline-MPC. (B) Mean number of rewards (spheres) by category. STM-MCTS collects more green and yellow (positive) rewards than the baseline agents. However, Baseline-MPC collects fewer (negative) red rewards, which is likely related its ability evaluate actions after every step. Uncertainty regions and bars indicate one standard deviation over runs.
2020 Crosby 2020 In AAI, an agent is tasked with reaching green sphere given particular setup that may include intermediary rewards (yellow spheres), terminal negative rewards (red spheres), obstacles (e-g. walls), etc. For the purposes of this work, we use a sparsely populated configuration with single green, red, and yellow spheres, in which successful agent would be forced to perform both short- and long-distane planning, as well as more extensive exploration of the environment.
5.1 EXPERIMENTAL RESULTS
We tested the STM agent using 100,000 steps in undomly-generater environments (max episode length of 500) against the baseline system with two different planning procedures MCTS and model-predictive control (MPC). In contrast to MCTS, the MPC agent re-evaluates its plan after every action. Figure 3 summarises the experimental results. Our STM-MCTS agent outperforms the baseline systems in acquiring more rewards within the 100,000 steps. In particular, we note that the STM-MCTS agent showed significant improvement against the Baseline-MCTS Similarly, we show that STM-MCTS model retrieves more cumulative reward than the Baseline-MPe agent, which uses computationally expensive planning procedure. Specifically, our agent achieves higher cumulative reward in less than half the time, ~-6 hours, compared to 14 hours.
5.2 ROLL-OUT INSPECTION
Inspecting prediction roll-outs produced by the STM-based system provides great insight into its practical benefits for the agent's performance. Specifically, agent is capable of varying the temporal extent of its predictions and imagining future salient events.
5.2.1 VARYING PREDICTION TIMESCALE
Much like human perception of time changes depending on the perceptual content of the surround- ings, our agent varies the prediction timescale depending on the context it finds itself in. Specif- ically, in the AAI environment the complexity of any given observation is primarily driven by the presence of objects, which may appear different sizes, colours, and configurations. As result, our agent consistently predicts farther into the future in the absence of any nearby objects, and slows its timescale, predicting at finer temporal rate, when the objects are close.
Practically, this has several important implications. First, performing temporal jumps and skipping unnecessary intermediary steps affords greater computational efficiency, and reduces the detrimental effects of sequential error accumulation, as can be seen in Figure Second, while STM is able to predict far ahead, its inherent flexibility to predict over varying timescales does not compromise the agent's 's performance when the states of interest are close. Thus, separate mechanism for adjusting how far into the future an agent should plan is not necessary and implicitly handled by our model. Third, STM allows the agent to make more informed decisions in environment, it tends to populate the roll-outs with salient observations of the short- and long- tergterm futures depending on the context. As result, STM effectively re-focuses the central purpose of transition model from

---

Under review as a conference paper at ICLR 2021
most accurately modeling the ground-truth dynamics of environment to predicting states more informative with respect to the affordances of the environment.
Roll-out Steps
Action: forward
Ground-truth
sequence
Objective-Timescal
STM
fast timescale
tinescale
Figure 4: STM vary prediction timescale, with large gaps between predictions at the start, which becomes more fine-grained as the object gets closer. timessale model suffers from slow-timescale predictions and error accumulation, resulting in poorly-informaiiee predictions.
5.2.2 IMAGINING SURPRISING EVENTS
As mentioned, S-sequences frequently include epistemically-suupiiii transitions, which, in the context of the AAI environment, constitute events where objects come into view. As a result, STM significantly
likely to include roll-outs with objects appearing in the frame, in contrast to the baseline that employs the objective-timescale transition model.
The ability of the STM to imagine novel and salient future events encourages exploratory behaviour, which is distinct from the active inference agent's intrinsic exploratory motivations. We again stress that although the predicted futures may be inaccurate with respect to the ground-truth positions of the objects, they are nevertheless more informative with respect to the agent's potential affordances in the environment. This in stark contrast with the objective-timestt futures in the absence of any objects. As result, the STM agent is less prone to get stuck in sub-optimal state, which was commonly observed in the baseline system, and more inclined to explore the environment beyond its current position.
model, which imagines
Roll-out Steps
Action: turn left
Objective-iimeccae
Model
Figure 5: STM is able to imagine surprising events. Despite the fact that appearance objects is rare event in the environment, STM frequently predicts them in the roll-outs. In contrast, objective- timescale model is not capable of that as a direct corollary its training procedure.
CONCLUSION AND FUTURE WORK
We proposed STM, a novel approach to learning a transition dynamics model with the use of se- quences of episodic memories, which define an agent's more useful, subjective timescale. STM showed significant improvement against the baseline agent's 's performance in the AAI environment. Inspired by the problems of inaccurate and inefficient long-term predictions model-based RL and the recent neuroscience literature on episodic memory and human time perception, we merged ideas from the different fields into one new technique of learning a forward model. We further emphasised two important characteristics of the newly-devised model its ability vary the temporal extent of future predictions and to predict future salient events. The application our technique not limited to active inference, and be adapted for use in other model-basee frameworks.
Future work may explore more generalised approaches of action summarisation and dynamic thresh- olding for memory formation. Another enticing direction of research is to investigate the feasibility of having a single transition model that slowly transitions from training on an objective timescale to training subjective timescale, the memory formation goes on.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, et al. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https:// WWw. www.tenscrilow.com org/ Software available from tensorflow.org.
Pieter Abbeel, Morgan Quigley, and Andrew Y. Ng. Using inaccurate models in reinforcement learning. In ICML '06, 2006.
J. Andrew Bagnell and Jeff G. Schneider. Autonomous helicopter control using reinforcement learn- ing policy search methods. Proceedings 2001 ICRA. JEEE International Conference on Roboties and Automation (Cat. No.01CH37164), 2: 10-161151622 vol.2, 2001.
Benjamin Beyret, José Hemandez--Orlloo Lucy G Cheke, Marta Halina, Murray Shanahan, and Matthew Crosby. The Animal-AI Environment: Training and testing animal-like artificial cogni- tion. ArXiv, abs/1909.07488, 2019.
Charles Blundell, Julien Corebiiee, K. Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. ArXiv, 155054405 2015.
Charles Blundell, Benigno Uria, Alexander Pritzel, Yazhe Li, Avraham Ruderman, Joel Z. Leibo, Jack W. Rae, Daan Wierstra, and Demis Hassabis. Model-free episodic control. ArXiv, abs/1600.04ee 04060, 2016.
Matthew M Borvinick, Sam Ritter, Jane X. Wang, Zeb Kurth-Nelson, and Demis Hassabis. Rein- forcement learning, fast and slow. Trends in cognilive sciences, 23 5:408-422, 2019.
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. In NeurlPS, 2018.
Catalin V. Buhusi and Warren H. Meck. What makes us tick? functional and neural mechanisms of interval timing. Nature Reviews Neuroscience, 6:755-765, 2005.
L. Cheke and N. Clayton. Eurasian jays (garrulus glandarius) overcome their current desires to anticipate two distinct future needs and plan for them appropriately. Biology Letters, 8:171 175, 2011.
Ignasi Clavera, Yao Fu, and Pieter Abbeel. Model-augmented actor-critic: Backpropagating through paths. ArXiv, abs/2005.08005 08068, 2020.
Nicola S. Clayton, Timothy Bussey, and Anthony Dickinson. Can animals recall the past and plan for the future? Nature Reviews Neuroscience, 4:685-691, 2003. 2003.
Matthew Crosby. Building thinking machines by solving animal cognition tasks. Minds and Ma- chines, 2020.
Matthew Crosby, Benjamin Beyret, Murray Shanahan, José Hermández-Orallo, Lucy Cheke, and Marta Halina. The Animal-AI testbed and competition. volume 123 of Proceedings of Machine Learning Research, pp. 164-176. PMLR, 2020.
Mare Peter Deisenroth and Carl E. Rasmussen. Pilco: model-based and data-efficient approach to policy search. In ICML, 2011.
Stefan Depeweg, José Miguel Hernndeez-Lobtoo Finale Doshi-Veler and Steffen Udluft. Learn- ing and policy search stochastic dynamical systems with bayesian neural networks. ArXiv, abs/1605.07122 2017.
Ben Deverett, Ryan Faulkner, Meire Fortunato, Greg Wayne, and Joel Z. Leibo. Interval timing in deep reinforcement learning agents. ArXiv, 0bb/1995511449) 2019.
Lasse Espeholt, Hubert Soyer, Rémi Munos, Karen Simonyan, Volodymyr Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, Shane Legg, and Koray Kavukcuoglu. Im- pala: Scalable distributed deep-rl with importance weighted actor-learner architectures. ArXiv, abs/1 fbs/1802.0156) 2018.

---

Under review as a conference paper at ICLR 2021
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I. Jordan, Joseph E. Gonzalez, and Sergey Levine. Model-based value expansion for efficient model-free reinforeement learning. 2018.
Zafeirios Fountas, Noor Sajid, Pedro AM Mediano, and Karl Friston. Deep active inference agents using monte-carlo methods. Accepted to NeurIPS, 2020a.
Zafeirios Fountas, Anastasia Sylaidi, Kyriacos Nikiforou, Anil K. Seth, Murray Shanahan, and War- rick Roseboom. A predictive processing model of episodic memory and time perception. bioRxiv, 2020b.
Karl J. Friston. free energy principle for particular physics. arXiv: Neurons and Cognition, 2019.
Karl J. Friston, James Kilner, and Lee Harrison. A free energy principle for the brain. Journal of Physiology-Paie Paris, 100:70-88 2006.
Karl J. Friston, Francesco Rigoli, Dimitri Ognibene, Christoph Mathys, Thomas H. B. FitzGerald, and Giovanni Pezzulo. Active inference and epistemic value. Cognitive Neuroscience, 6:187 214, 2015.
Karl J. Friston, Thomas H. B. FitzGerald, F. Rigoli, Schwattenbeck, O' Doherty, and G. Pezzulo. Active inference and learning. Neuroscience and Biobehavioral Reviews, 68:862-877 877, 2016.
Karl J. Friston, Thomas H. B. FitzGerald, F. Rigoli, Schwartenbeck, and G. Pezzulo. Active inference: A process theory. Neural Computation, 29:1 2017a.
Karl J. Friston, Thomas H. B. FitzGerald, Francesco Rigoli, Philipp Schwartenbeck, and Giovanni Pezzulo. Active inference: A process theory. Neural Computation, 29:1-49) 2017b.
Karl Friston, Marco Lin, Christopher Frith, Giovanni Pezzulo, Allan Hobson, and Sasha Ondobaka. Active inference, curiosity and insight. Neural Computation, 29:2633-2683 2017c
Karl J. Friston, Richard E Rosch, Thomas Parr, Cathy J. Price, and Howard Bowman. Deep temporal models and active inference. Neuroscience and Biobehavioral Reviews, 90:486 -501, 2017d 7d.
Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncertainty in deep learning. In ICML, 2016.
Samuel J. Gershman and Nathaniel D. Daw. Reinforcement learning and episodic memory in hu-
mans and animals: An integrative framework, Annual Review of Psychology, 68:101-128, 2017.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. ArXiv, abs/1410.5401,
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, et al. Hybrid computing using neural network with dynamic external memory. Nature, 538:471-476.1 2016.
Daniel L. Greenberg and Mieke Verfaellie, Interdependence of episodic and semantic memory: evidence from neuropsychology. Journal of the International Neuropsycholgical Society : JINS, 16 5:748-53, 2010.
Andrea Greve, Elisa Cooper, Alexander Kaula, Michael C. Anderson, and Richard N. A. Henson. Does prediction error drive one-shot declarative learning? Journal of Memory and Language, 94: 149 165, 2017.
David R Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In NeurIPS, 2018.
Jessica B. Hamrick. Analogues of mental simulation and imagination in deep learning. Current Opinion in Behavioral Sciences. 29:8-16, 2019.
Steven Hansen, Pablo Sprechmann, Alexander Pritzel, Andr'e e Barreto, and Charles Blundell. Fast deep reinforcement learning using online adjustments from the past. In NeurIPS, 2018.
10

---

Under review as a conference paper at ICLR 2021
Demis Hassabis, Dharshan Kumaran, Seralynne D. Vann, and Eleanor A. Maguire. Patients with hippocampal amnesia cannot imagine new experiences. Proceedings of the National Academy of Sciences, 104:1726 1731, 2007.
Thomas T. Hills. Towards unified theory of animal event timing. 2003.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9:1735- 1780, 1997.
Chia-Chun Hung, T. Lillicrap. Josh Abramson, Yan Wu, M. Mirza, F. Carnevale, Arun Ahuja, and G. Wayne. Optimizing agent behavior over long time scales by transporting value. Nature Com- munications, 10, 2019.
Anthony Jang, Matthew Nassar. Daniel G. Dillon, and Michael Frank. Positive reward prediction errors strengthen incidental memory encoding. bioRxiv, 2018.
Hyunwoo Jung, Moonsu Han, Minki Kang, and Sung Ju Hwang. Learning what to remember: Long-term episodic memory networks for learning from streaming data. ArXiv, abs/1812.04227, 2018.
Lukasz Kaiser, Ofir Nachum, Aurko Roy, and Samy Bengio. Learning to remember rare events.
ArXiv, abs/1703.03129, 2017.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model-based reinforcement learning for atari. ArXiv, abs/1903.0000 00374, 2020.
Gabriel Kalweit and Joschka Boedecker. Uncertainty-driver imagination for continuous deep rein- forcement learning. In CoRL, 2017.
Ken Kansky, Tom Silver, David A. Mély, Mohamed Eldawy, Miguel Lázaro-Gredilla Xinghua Lou, N. Dorfman, Szymon Sidor, Scott Phoenix, and Dileep George. Schema networks: Zero-shot transfer with generative causal model of intuitive physics. ArXiv, Abs/1770.04337,, 2017.
Nan Rosemary Ke, Amanpreet Singh, Ahmed Touati, Anirudh Goyal, Yoshua Bengio, Devi Parikh, and Dhruv Batra. Learning dynamics model in reinforcement learning by incorporating the long term future. ArXiv, abs/ 0tbs//903.01999 2019.
Jonathan Ko, Daniel J. Klein, Dieter Fox, and Dirk Hähnel. Gaussian processes and reinforcement learning for identification and control of an autonomous blimp. Proceedings 2007 IEEE Interna- tional Conference on Robotics and Automation, pp. 742-747,2000
Vikash Kumar, Emanuel Todorov, and Sergey Levine. Optimal control with leamned local models: Application to dexterous manipulation. 2016 IEEE International Conference on Robotics and Automation (ICRA), pp. 378-33-- 778-383,22 2016.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. ArXiv, abs/1802..1092. 2018.
Máté Lengyel and Peter Dayan. Hippocampal contributions to control: The third way. In NIPS, 2007.
Sergey Levine and Pieter Abbeel. Learning neural network policies with guided policy search under unknown dynamics. In NIPS, 20
Sergey Levine and Pieter Abbeel. Learning neural network policies with guided policy search under unknown dynamics, In NIPS, 2014b.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuo- motor policies. J. Mach. Learn. Res., 17:39:1-39:40, 2016.
Johannes B. Mahr and Gergely Csibra. Why do we remember? the communicative function of episodic memory. The Behavioral and brain sciences, pp. 1-93, 2017.

---

Under review as a conference paper at ICLR 2021
Rowan McAllister and Carl Edward Rasmussen. Improving pilco with bayesian neural network dynamics models. 2016.
Kourken Michaelian, Mental time travel: Episodic memory and our knowledge of the personal past. 2016.
Beren Millidge. Deep active inference as variational policy gradients. ArXiv, 000s/1907.0335 2019.
Nikhil Mishra, Pieter Abbeel, and Igor Mordatch. Prediction and control with temporal segment models. ArXiv, abs/ 700704070000 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin A. Riedmiller, et Human-level control through deep reinforcement learning. Nature, 518:529-533. 2015.
Anusha Nagabandi, Gregory Kahn, Ronald S. Fearing, and Sergey Levine. Neural network dy- namics for model- based deep reinforcement learning with model-free fine-tuning. 2018 IEEE International Conference Robotics and Automation (ICRA). pp- 7559-7566, 2018.
Junhyuk Oh, Valliappa Chockalingam, Satinder P. Singh, and Honglak Lee. Control of memory, active perception, and action in minecraft. ArXiv, abs/1605.09 02188, 2016.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adrià Puigdomènech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. ArXiv,
abs/1703.,,988, 2017.
Sébastien Racanière, Theophane Weber, David P. Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, et Imagination-augmenntd agents for deep reinforcement learning- In NIPS, 2017.
Warrick Roseboom. Z. Fountas, Kyriacos Nikiforou, David Bhowmik, M. Shanahan, and A. Seth. Activity in perceptual classification networks as a basis for human subjective time perception. Nature Communications, 10, 2019.
Nina Rouhani, Kenneth A. Norman, and Yael Niv. Dissociable effects of surprising rewards on learning and memory. Journal of Experimental Psychology: Learming, Memory, and Cognition, 44:1430-14453 2018
Noor Sajid, Philip Ball, and Karl Friston. Active inference: demystified and compared. arXiv: Artificial Intelligence, 2019.
D. Schacter, D. Addis, D. Hassabis, V. C. Martín, R. N. Spreng, and K. Szpunar. The future of memory: Remembering, imagining, and the brain. Neuron, 76:677-694, 2012.
Daniel L. Schacter, Donna Rose Addis, and Randy L. Buckner. Remembering the past to imagine the future: the prospective brain. Nature Reviews Neuroscience, 8:657-661, 2007.
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay.
CoRR, abs/1511.0500 .05952, 2016.
Maxine Sherman, Zafeirios Fountas, Anil K Seth, and Warrick Roseboom. Accumulation of salient perceptual events predicts subjective time. bioRxiv, 2020.
D. Silver, Julian Schritwieser, K. Simonyan, Ioannis Antonoglou, Aja Huang, A. Guez, Hubert, L. Baker, et al. Mastering the game of go without human knowledge. Nature, 550:354-353 2017.
Alexander Tschantz, Manuel Baltieri, Anil K. Seth, and Christopher L. Buckley. Scaling active inference. ArXiv, abs/1911. 10601,, 2019.
Endel Tulving. Episodic and semantic memory. 1972.
Kai Ueltzhöffer. Deep active inference. Biological Cybernetics, 112:547-573, 2018.
Manuel Watter, Jost Tobias Springenberg. Joschka Boedecker, and Martin A. Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In NIPS, 2015.
12

---

Under review as a conference paper at ICLR 2021
Guangxiang Zhu, Zichuan Lin, Guangwen Yang, and Chongjie Zhang. Episodic reinforcement learning with associative memory. In ICLR, 2020.
Ozan Çatal, Johannes Nauta, Tim Verbelen, Pieter Simoens, and Bart Dhoedt. Bayesian policy selection using active inference. ArXiv, abs/1904.08144 2019.
13

---

Under review as a conference paper at ICLR 2021
A PRELIMINARIES
A.1 ACTIVE INFERENCE
Active inference is a corollary of the free-energy principle applied to action (Friston et 2016 Friston 2019 Sajid et al. 2019). In this framework, an agent embedded in an environment aims to do two things: (i) minimise surprisal from the observations the environment under the agent's internal model of this environment, and (ii) perform actions so to minimise the expected surprisal in the future. More formally, an agent is equipped with a generative model p(Or, S 0), where or is the agent's observation at time S, is the hidden state of the environment, and 0 denotes the parameters of the generative model. The agent's ' surprise at time is defined as the negative log- likelihood. - log p(o;; 0).
We can upper-bound this intractable expression using variational inference by introducing an ap- proximate posterior distribution, (St)), over St, such that:
- -logg(o; 0) < (s((s)) [log ((s)) - log p(o, S;; 0)] = F,
where FF the variational free energy. The minimisation of this quantity realises objective (i) and is performed by optimising the parameters of the generative model, 0. It is also equivalent to the maximisation of model evidence, which intuitively implies that the agent aims to perfect its generative model explaining the sensory observations from the environment. To realise objective (ii), the agent must select actions that lead to the lowest expected surprise in the future, which can be calculated using the expected free energy (EFE), G:
G(T,T) r) = Ep(o-1+-) (9((+(x) [log g( (T) - log plor s-1T+]]
(8)
variational free energy, F
where т L and {ar, Ge+11 1; ay } is sequence of actions (policy) between the present time and the future time r. The free-energy minimising system must, therefore, imagine the future observations given a policy and calculate the expected free energy conditioned on taking this policy. Then, actions that led to lower values of the EFE are chosen with higher probability. as opposed to actions that led to higher values of EFE, such that:
p(T) a (-)G(T)),
where G(T)= = -E., G(T,, r), is the temperature parameter, o-)) denotes a softmax function, and t is the present timestep.
B ARCHITECTURAL DETAILS AND TRAINING
B.1 BASELINE IMPLEMENTATION
As mentioned, each component of the generative and inference models is parametrised by feed- forward neural networks (including fully-connected, convolutional and causpose-convolutione lay- ers), whose architectural details can be found in Figure 6 The latent bottleneck of the autoencoder, s. was of size 10. The hyperparameters of the top-down attention mechanism were: 2, 0.5, 0.1, and d = 5, chosen match those in Fountas et al. (2020a). Similarly, we restricted the ac- tion space to just actions - forward, left, right. For testing, we optimised the MCTS parameters of the baseline agent, setting the exploration hyperparameter Gesplore 0.1 (see Eq 10). T0 and performing 30 simulation loops, each with depth of 1. The networks were trained using separate optimisers for stability reasons. The habitual and transition networks trained with learning rate of 0.0001; the autoencoder's s optimiser had a learning rate of 0.001. The batch size was set to 50 and the model was trained for 750k iterations under green observational prior. All of the networks were implemented using Tensorflow v2.2 (Abadi et al. 2015). Tests were performed in Animal-AI v (Beyret et al. 2019).
Furthermore, following Fountas et al. 2020a), we define the MCTS upper confidence bound as,
as, Ō(s a) Cexplore Qo (als). N(s, a) +1
(10)
14

---

Under review as a conference paper at ICLR 2021
q( (tri xba)
P( 1s, 4r; 0,)
( Φ,)
p(or 1s, 0s)
Habitual Network
Transition Model
Encoder
Decoder
Iepuz shope: [10.)
Input shope: (10+): (10+33
Inpur Inpurshape: shape: (32, 32, 3)
Enpat shape: (10)
Com2DD: 32 (3, stride
Linear: untt
ReL.U
ReL.U
Rel.U
Dropout 0.5
Dropout
Lineer: 256 mits
Lins:: 128
Linea: 512 ReL.U
Com2DD 54 (3, stride
ReLU
ReLU
Dropout
Dropout 0.5
Linrar: units
Lineor: 512 unit RaLU
RcLU
Softmax
Com22D: 128 filien krezel 3), Ret.U
Drapout 0.5
stride (2, 2)
Dropout 0.5
Lineer: 16-15*64
Rel.U
Outpur shaper
Libea: 512 uilts ReL.U
Dispout
Com2D: 256 fillen, knezel (3,
Con2DTrassees:
Drapout 0.5
2)
256 fliters kemel (3, 3), strde UI. SAME padking, RelU
Liteer: 20 units No Autivation
I Linean 256 unis Ret.U
ComcnTTansper
Ousput shape: (20))
125 filters, kernel 3) stide (2. SAME pabking,
Dropout
Limar: 256 unls Ret.U
Dropout 0.5
kemel aride
SAME parking, ReL.U
Linear: 256 uniss Rel.U)
Dropost U.5
3 Sikter, karnel (3, 3), suide (1, 1) SAMI pwtding Sigmoit
Linoer: Activation
Ostput shope: (20)
Outpot shope (30,32, 3)
Figure 6: Implementation of the baseline system.
As discussed, each network was trained with its corresponding loss function, which are the con- stituent parts of the total variational free energy. In particular, the autoencoder was trained using Eqs. Ta and Tb transition model using Eq. [b and habitual network using Eq. Ic
Furthermore, following the training procedure from Fountas et al. (2020a), we stabilise the conver- gence of the autoencoder by modifying the loss function to:
Cautoencoder = E [logp(o, s,; 0,)l + DDi.. o. 1
(11)
+ (1- (q(s;; (0, I)],
where y is hyperparameter that gradually increases from 0 to 0.8 during training.
B.2 PRIORITISED EXPERIENCE REPLAY
As part of the baseline system's 's training procedure, we utilise prioritised experience replay (PER) (Schaul al. 2016) to mitigate the detrimental effects of on-line learning (which was used in the original paper by Fountas et al. (2020a)). and to encourage better object-centric representations.
In particular, on-line learning has three major issues associated with it. First, training is performed on correlated data points, which is generally considered to be detrimental for training neural net- works (Schaul et al. (0166). Second, observations that are rarely encountered in an environment are discarded in on-line learning and are used for training only when visited again. These are likely to be the observations for which there is most room for improvement. Instead, the agent will often be training already well- rell-predicted transitions that it happens to visit often. Finally, on-line learning agent is constrained by its current position in an environment to sample new data and, thus, has very limited control over the content of its training batches.
15

---

Under review as a conference paper at ICLR 2021
Furthermore, mentioned in Section in the Animal-AI enviromeent rare observations are those that include objects; yet, objects are a central component of this environment the only way to interact, get rewards, and importantly the only means of minimising the free energy optimally. To encourage our agent to learn better object representations, we employ PER with the objective- timescale transition model free energy the priority metric. As discussed, observations with higher values of this metric tend to constitute more complex scenes, which include objects as the only source of complexity in the AAI. See Figure 7 for qualitative evidence this trend. The use PER resulted in a considerable improvement in the baseline's 's performance and better ability to reconstruct observations with objects (See Figure
B.3 STM IMPLEMENTATION
The STM introduces two additional components: STM habitual network and STM transition model. The habitual network was trained using the same training procedure as described in Appendix [ The transition model was trained batch size 15 and learning rate of 0.0005. Each batch consisted of zero-padded S-sequences with length 50. We use a Masking layer to ignore zero-padded parts of the sequences in the computational graph. The training was stopped at 200k training iterations. For testing STM-MCTS in Section [ we optimise the MCTS parameters, setting Gexpiore = 0.1, and performing 15 simulation loops, each with depth of The threshold (objective-timescal transition model free energy), was manually set to 5 after inspection of the buffer and value distribution.
B.4 ACTION HEURISTIC
To train the STM transition model in the Animal-AI environment, we implement a simple heuristic that is used summarise sequence of actions taken by the agent from one memory to reach takes the agent from a
the next one. A sequence of actions, A = (aT, +1,-4++++ recorded memory 8r, memory 8yy, where the time between these states 72 Ti N, and E {Aforward; Oright; Glef). ). We employ polar coordinates relative to the agent's initial position in Cartesian coordinates time Ti and perform iterative updates of its position after every action un- til the time-step of the next episodic memory, is reached. Given the agent's 's orientation in the environment, 6, the next position of the agent calculated using,
sin A cos 0]
Pi+1 = Ps
where Pt
(12)
t=
Finally, we retrieve angle o, which describes the direction which the agent has travelled with re- spect to its initial position and orientation. This angle is used decide the action that summarises the trajectory using
Oforward
|0| S 22.5
Oright
22.5 < Φ 180
Gleft
22.55 2 180%,
Although this heuristic provided satisfactory results, trajectory encoding one of the most limiting parts of the STM and is promising direction for further research.
16

---

Under review as a conference paper at ICLR 2021
More objects
Highest
Median
Fewer objects
Lowest
Figure 7: Observations sorted by their corresponding recorded value of the objective-timescale tran- sition model free energy in descending order. States with higher values on average contain more objects, constituting complex settings.
17

---

Under review as a conference paper at ICLR 2021
A
B
Roll-out Steps
Action: forward
On-line
PER
Ground-truth
On-line Baseline
PER Baseline
Ground-truth
Action: turn left
On-line Baseline
PER Baseline
Ground-trute
Figure 8: Agent trained with PER is better at generating observations with objects. (A) Prediction roll-outs showing that baseline trained with PER can better represent objects, thus preserving them in future predictions. (B) Ground-truth observations (third column) are passed through the autoencoder of the on-line and PER agents. As can be seen, observations are better reconstructed by the baseline system trained with PER.
q(a, |s da) STM Habitual Network
p(s (Se. ar, 1x=1):
STM Transition Model
Input shape: (10.)
Inpur shape:
Linear: 120 units
LSTM: 512 units Rel.1
Sor (Stx der hz-1)
Linear: 128 units Rel.U
LSTM: 512 mils
} p(s |h, Ons)
Linear 128 units ReLU
Linear: 20 units Activation
Drogour
Lioea:: ReL.U
Ourpur
20.)
Linear: units Soltmax
Oorput shape: (3)
Figure 9: Implementation of STM components.
18

---

Under review as a conference paper at ICLR 2021
C ADDITIONAL RESULTS
We provide additional results of the STM prediction roll-outs:
a) Figure 10: random roll-outs generated by the system. These diverse roll-outs demonstrate that STM able to: i) make correct aetion-conditioned predictions, ii) speed up pre- diction timescale when objects are far away, iii) slow down the prediction timescale when objects are nearby.
b) Figure 11 STM consistently imagines objects coming into view. The observations pro- duced by the model are entirely plausible given the path the agent is taking and the context it finds itself predictions. It is pertinent to note that the roll-outs comply with the physics of the environ- ment, which is crucial, as potentially refutes the hypothesis that these imagined objects were predicted at random.
This indicates that STM does indeed produce semantically meaningful
c) Figure [2] shows the roll-outs produced by the senective-timescale model using the same starting states as in Figure [I] These roll-outs are in stark contrast to those produced by STM, exemplifying the baseline's inability to imagine objects that are not present in the initial frame.
19

---

Under review as a conference paper at ICLR 2021
Starting State
Imagined Roll-out
Forward
Left
Right
Forward
Left
Right
Forward
Left
Right
Forward
Left
Right
Forward
Left
Right
Forward
Left
Right
F2
Figure 10: Random roll-outs generated with the STM transition model.
20

---

Under review as a conference paper at ICLR 2021
Starting State
Imagined Roll-out
Forward
Left
Right
Forward
Left
Right
Forward
Left
Right
Forward
Left
Right
Forward
Lefi
Right
Forward
Left
Right
T2
F4
85
F6
F7
T9
Figure 11: STM can imagine objects from 'uninteresting" states, Arrows indicate roll-outs with imagined objects.
21

---

Under review as a conference paper at ICLR 2021
Starting State
Imagined Roll-out
Forward
Left
Right
Forward
Left
Right
Forward
Left
Right
Forward
Left
Right
Forward
Left
Right
Forward
Right
t=0
t=1
t=2
t=3
1=4
1=5
C-6
1=7
[=8
t=9
Figure 12: contrast to STM, the objective objective-timesaal timescale transition model is not able to imagine ob- jects, starting with the same initial observations shown in Figure II
22


Reviews:
[{"review_id": "ksmtpkZvju7", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nMost model-based RL algorithms learn dynamics models that predicts the next timestep. However, because of model-bias, frequency of timesteps, and objective timescales, the dynamics models can accumulate errors and limited by timescales. The authors propose subjective-timescale model (STM) that instead of predicting the next timesteps they find the \"surprising\" subsequences of the trajectories and learn temporal-skipping dynamics models over them. The paper shows the improvement over single-step prediction baselines in a first-person navigation domain.\n\nPros:\n\nThe method aims to address a very important problem in model-based RL. \n\nThe idea of using variational free energy with model-based RL seems novel to me, and has not been widely explored.\n\nThe qualitative visualization (figure 4 and 5) provides a nice understanding of what the method is doing as well as what it is capable of in first-person navigation. \n\nCons:\n\n-- Methods --\n\nThe main sections do not contain sufficient information regarding how the actions are obtained from learned STM models. I find one paragraph in section 3 in sufficient. The experimental sections also do not mention how MCTS and MPC baselines differ. Please clarify. Also, how do you recover low-level action sequences from the aggregated actions after MCTS? I do not find the answer from the paper.\n\nThe keyframe selection method requires more justification. It is unclear how using the KL divergence for measuring surprise will improve over model's prediction error studied in previous work.\nIt would be amazing to provide a theoretical justification for the heurstics toward, e.g., saying something about the end task performance, if possible. How sensitive is it to the KL threshold? Please provide this study.\n\nThe action sequence aggregation is domain specific which seems a bit unfair to compare against the baselines which **do not** have access to the same information. There should be more baselines or ablation studies to disentangle the improvement of the method from this domain-specific assumption.\n\nThe paper uses different indexing styles which make the method more confusing than it should have been. Please choose one between indexing tau or arithematics on tau.\n\n-- Experiments --\n\nAs metioned briefly before, more baselines or ablation will be critical to judge the importance of the proposed model? What about compare against other sliency approaches such as prediction error for memory accumulation? Also, it would be helpful to have more than 1 environment to show the genrality of the approach.\n\nThe experimental results do not provide enough information to understand what tasks can be solved and what cannot be solved in Animal AI environment. Can the we provide the success rates and categorized by difficulty levels? These information will be helpful in understanding what STM can and cannot do. Perhaps, having a link to some videos will also help.\n\nFigure 3 shows that the reward is going up; how far can the rward go? It is still increasing. \n\nIn figure 3, how is the reward computed?\n\nIn figure 5, is there a groundturth trajectory comparison?\n\n\n-- Others --\n\nPersonally, I would appreciate more background and intuition on each term in the variational free energy formula.\n\nConclusion:\n\nAgain, I believe that this work is addressing a very important question with an interesting idea, but it may require a little bit more work to make the case. I appreciate the authors thinking about this problem, and hope the authors are encouraged to continue their work.", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ya2jyJegbdK", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: The authors propose to train a model not on the objective time-scale of a sequence of frames, but on the subjective time-scale dictated by how surprising events are (where surprise here is defined as being above a certain energy free threshold). The trained action-conditioned model learns to slow down time for complex scenes, and fast forward when things are easily predicted.\n\nThe overall topic is an important one, most model based methods suffer from accumulating errors.\nThe introduction is well written and offers a strong motivation for the rest of the paper. I like the explanation in terms of a distinction between objective and subjective perception of time and events.\n\nThere are lots and lots of references to work by Friston et al., but I am not sure I see the connection with active inference as being that strong or even necessary.\n\nRegarding the part “Furthermore, for long-term predictions STM systematically performs temporal jumps (skipping intermediary steps), thus providing more informative future predictions and reducing the detrimental effects of one-step prediction error accumulation.” There are two potentially relevant references (Neitz et al NeurIPS  2018), and (Darajaman et al, ICML 2019), as both learn models (not necessarily action-conditioned though) that can skip an adaptive number of steps into the future, with similar consequences (i.e. preventing error accumulation and increasing rollout speed).\n\nP4: “The habitual network acts as a model-free component of the system, learning to map inferred states directly to actions”. Shouldn’t $q(a_t; φ_a)$  be a function of $s$ as well?\n\nP6: The angle heuristic might deserve a more in-depth discussion. What if the agent goes in a circle? Or does a U turn? Or a complex sequence of movements in a maze? \nSince the goal of STMs is (at least in part) to reduce progressively the length of S-sequences (such that they start spanning longer and longer horizons), the actions that bridge two episodic memory need to be summarised in a way that is expressive enough.\nCould one think of applying the same STM principle that is already applied to states, to actions as well?\n\n“Importantly, [the] function $f_{\\theta_{s}}$ is deterministic and serves only to encode information about preceding…”. Where does $f_{\\theta_{s}}$ appear in this context? I can’t see it anywhere between eq. 4 and 6\n\nThis AAI environment is not so established. Perhaps the author could add (even in the appendix) a top view of a typical setting that corresponds to one episode? \nAlso, it appears that the setting chosen by the authors does not reflect the characteristics that they planned to showcase with their method. Given the presence of one sphere per colour and a relatively small environment, how can we appreciate that the agent learns to do planning over long horizons?\nPerhaps an environment with longer horizons would be a more adequate testbed?\n\nI very much liked that the authors showed the additional results in Appendix C, I think they are extremely important for the paper. However, I disagree with the claims made in the text, as it appears that they are not substantiated by the figures they reference.\n- “Figure 10: random roll-outs generated by the system. These diverse roll-outs demonstrate that STM is able to: i) make correct action-conditioned predictions, ii) speed up its prediction timescale when objects are far away, iii) slow down the prediction timescale when objects are nearby”\nIf I interpret Figure 10 correctly, it seems to suggest that the imagined roll-outs are not very consistent, as in almost every single case the color of the spheres changes from green to yellow and vice-versa during the rollout.\n- Figure 11: several transitions appear not to be realistic, objects appear out of nowhere instead of smoothly while turning.\n- Figure 12: while it is true that the baseline model does not generate any sphere, it appears to me that the physics is significantly more consistent than with STM.\n\nThe last sentence in the conclusion seems to suggest that the model is not progressively expanding its horizon the more it trains. How come? \n\nOverall, my impression is that the paper has a very interesting motivation, but the execution could be significantly improved. The results are not so convincing, due to:\n1) figures (10-11-12) that do not soundly corroborate the claims, \n2) evaluation setting that does not allow to really test the claim (i.e. not enough opportunities for increasingly longer horizons as the model improves) \n3) the lack of established baselines and benchmarks\n4) the lack of an algorithm box to present the method\n\nMinor:\n- Simplify language where possible (utilised -> used; are capable of -> can; etc.)\n- Several figures (4,5, 8) are corrupted using preview on a Mac (not sure if it’s just my computer). I could see them correctly by using Chrome to open the PDF.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "AyGJlO-Vzsl", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this paper, the authors describe a variable-timescale prediction model for planning in the context of a deep active inference agent. They show that this agent outperforms a baseline in a scavenging task in a 3D first person environment. They show example rollouts of the baseline and variable-timescale models.\n\nUntil reading this paper I wasn't familiar with deep active inference agents, which apparently enable the extension of free energy methods to more complex settings. It seems like an intriguing alternative to deep RL. It's not clear to me whether it has the scaling potential that has been demonstrated for deep RL. Although the experiments reported here are larger scale than experiments with active inference systems prior to Fountas et al. (2020), they seem to be quite simple in comparison to tasks on which deep RL agents excel (e.g. typical Vizdoom and DMLab tasks). I don't mean this as a criticism, but instead as a question mark: this difference in demonstrated scale might simply be due to the massive difference in resources that have been devoted to deep RL agents vs active inference agents. \n\nAs a result, I tried to evaluate the paper on its merits in the context of active inference systems, instead of via comparison with deep RL. As such, I'd be happy with a convincing demonstration that (1) the variable timescale model outperforms a strong time-locked model in the scavenging task and (2) that its performance is due to the selection of useful frames for planning that MCTS then uses effectively.\n\nFigure 3 seems to answer (1) in the affirmative. I'm basically trusting the authors that theirs is a reasonably strong baseline, since I don't have experience training active inference agents or with this scavenging environment. This will be reflected in my confidence score.\n\nHowever, (2) is not substantiated well by the paper's analysis section. The only evidence for this is in the form of two pairs of example rollouts for the time-locked and variable-time models. Are these cherry picked examples, or are they actually reflective of general trends? I'd be much more comfortable if the authors supplied some aggregate results to substantiate this claim:\n\n\"As a result, our agent consistently predicts farther into the future in the absence of any nearby objects, and slows its timescale, predicting at finer temporal rate, when the objects are close.\" \n\nI'd really need aggregate results demonstrating that the S-sequences are summarizing long trajectories in a sensible way over a large set over episodes, to feel confident that the system is providing the benefits its purported to. Even better would be to analyze the MCTS search trees to show that the search trajectories over S-seqences have desirable properties.\n\nPerhaps more concerningly, the authors say that \"As a result, the STM agent is less prone to get stuck in a sub-optimal state, which was commonly observed in the baseline system, and is more inclined to explore the environment beyond its current position\". But, as far as I can tell, there's no evidence in the main paper for this. \n\nOne concrete concern is that the time-locked agent fails simply because it's rollouts aren't long enough to find the rewarding object. Maybe it would be sufficient to simply randomly drop out timesteps from the trajectory to reach the STM-MCTS performance level. The current analyses (and baseline results) don't rule out this hypothesis.\n\nIf the authors can provide stronger evidence on these points, I'd be very happy to increase my rating.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper uses a free-energy formulation to develop an approach to learning "jumpy" transition models, which predict surprising future states. This transition model is used in combination with MCTS and applied to a scavenging task in the Animal AI Olympics, outperforming two baselines.

While the reviewers praised the importance of the problem tackled, and the novelty of using a free energy approach, there was a general sense amongst the reviewers that the paper wasn't totally clear (especially for an RL audience). R1 also felt that some of the claims of the paper weren't sufficiently evaluated enough, and several reviewers indicated that they felt the baselines were insufficient (or, at a minimum, not described in enough detail to evaluate whether they were sufficient). Given these points, I feel the paper is not quite ready for publication at ICLR. I encourage the authors to flesh out their analysis a bit more, better describe the baselines (and possibly compare to other existing approaches as mentioned by R4), and overall to frame the paper a bit more for the RL community.

One additional reference the authors may be interested in: Gregor et al (2018). Temporal difference variational auto-encoder.
Decision: Reject