Original: Episodic Memory for Subjective-Timescale Models
Revised: Episodic Memory for Subjective-Timescale Models

Original: Anonymous Authors

Original: Abstract
Revised: Abstract

Original: Planning in complex environments requires reasoning over multi-step timescales.
Revised: Planning in complex environments requires reasoning over multi-step timescales.

Original: However, in model-based learning, an agent’s model is more commonly deﬁned over transitions between consecutive states.
Revised: However, in model-based learning, an agent’s model is more commonly deﬁned over transitions between consecutive states.

Original: This leads to plans using intermediate states that are either unnecessary, or worse, introduce cumulative prediction errors.
Revised: This leads to plans using intermediate states that are either unnecessary, or worse, introduce cumulative prediction errors.

Original: Inspired by the recent works on human time perception, we devise a novel approach for learning a transition dynamics model based on the sequences of episodic memories that deﬁne an agent’s subjective timescale – over which it learns world dynamics and over which future planning is performed.
Revised: Inspired by the recent works on human time perception, we devise a novel approach for learning a transition dynamics model based on the sequences of episodic memories that deﬁne an agent’s subjective timescale – over which it learns world dynamics and over which future planning is performed.

Original: We analyse the emergent beneﬁts of the subjective-timescale model (STM) by incorporating it into two disparate model-based methods –
Revised: We analyse the emergent beneﬁts of the subjective-timescale model (STM) by incorporating it into two disparate model-based methods –

Original: Dreamer and deep active inference.
Revised: Dreamer and deep active inference.

Original: Using 3D visual foraging tasks, we demonstrate that STM can systematically vary the temporal extent of its predictions and is more likely to predict future salient events (such as new objects coming into view).
Revised: Using 3D visual foraging tasks, we demonstrate that STM can systematically vary the temporal extent of its predictions and is more likely to predict future salient events (such as new objects coming into view).

Original: In comparison to the agents trained using objective timescales, STM agents also collect more rewards due to their ability to perform ﬂexible planning and a more pronounced exploratory behaviour.
Revised: In comparison to the agents trained using objective timescales, STM agents also collect more rewards due to their ability to perform ﬂexible planning and a more pronounced exploratory behaviour.

Original: 1. Introduction
Revised: 1. Introduction

Original: An agent endowed with a model of its environment has the ability to predict the consequences of its actions and perform planning into the future.
Revised: An agent endowed with a model of its environment has the ability to predict the consequences of its actions and perform planning into the future.

Original: Models allow agents to simulate the possible action-conditioned futures from their current state, even if the state was never visited during learning.
Revised: Models allow agents to simulate the possible action-conditioned futures from their current state, even if the state was never visited during learning.

Original: As a 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`ere et al., 2017; Mishra et al., 2017). The most popular framework for developing agents with in1 Anonymous Anonymous Country
Revised: As a 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

Original: [Figure]
Revised: [Figure]

Revised: (Racani`ere et al., 2017; Mishra et al., 2017).

Original: ternal models is model-based reinforcement learning (RL).
Revised: The most popular framework for developing agents with internal models is model-based reinforcement learning (RL).

Original: Model-based RL has seen great progress in recent years, with a number of proposed architectures attempting to improve both the quality and the usage of these models (Racani ` ere et al., 2017; Kansky et al., 2017; Hamrick, 2019;
Revised: Model-based RL has seen great progress in recent years, with a number of proposed architectures attempting to improve both the quality and the usage of these models (Racani ` ere et al., 2017; Kansky et al., 2017; Hamrick, 2019;

Original: Kaiser et al., 2020).
Revised: Kaiser et al., 2020).

Original: Nevertheless, learning internal models affords a number of unresolved problems.
Revised: Nevertheless, learning internal models affords a number of unresolved problems.

Original: 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, 2011).
Revised: 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, 2011).

Original: Long-term predictions are additionally computationally expensive in environments with slow temporal dynamics, given that all intermediary states must be predicted.
Revised: Long-term predictions are additionally computationally expensive in environments with slow temporal dynamics, given that all intermediary states must be predicted.

Original: Moreover, slow world dynamics 1 can inhibit the learning of dependencies between temporallydistant events, which can be crucial for environments with sparse rewards.
Revised: Moreover, slow world dynamics 1 can inhibit the learning of dependencies between temporallydistant events, which can be crucial for environments with sparse rewards.

Original: Finally, the temporal extent of future predictions is limited to the objective timescale of the environment over which the transition dynamics has been learned.
Revised: Finally, the temporal extent of future predictions is limited to the objective timescale of the environment over which the transition dynamics has been learned.

Original: This leaves little room for ﬂexible and context-dependent planning over varying timescales which is characteristic to animals and humans (Clayton et al., 2003; Buhusi & Meck, 2005; Cheke & Clayton, 2011).
Revised: This leaves little room for ﬂexible and context-dependent planning over varying timescales which is characteristic to animals and humans (Clayton et al., 2003; Buhusi & Meck, 2005; Cheke & Clayton, 2011).

Original: The ﬁnal issue exempliﬁes the disadvantage of the classical view on internal models, in which they are expected to capture the ground-truth transition dynamics of the environment.
Revised: The ﬁnal issue exempliﬁes the disadvantage of the classical view on internal models, in which they are expected to capture the ground-truth transition dynamics of the environment.

Original: Furthermore, in more complex environments with ﬁrst-person observations, this perspective does not take into account the apparent subjectivity of ﬁrst-person experiences.
Revised: Furthermore, in more complex environments with ﬁrst-person observations, this perspective does not take into account the apparent subjectivity of ﬁrst-person experiences.

Original: In particular, the agent’s learned representations of the environment’s transition dynamics implicitly include information about time .
Revised: In particular, the agent’s learned representations of the environment’s transition dynamics implicitly include information about time .

Original: Little work has been done to address the concept of time perception in model-based agents.
Revised: Little work has been done to address the concept of time perception in model-based agents.

Original: Empirical evidence from the studies of human and animal cognition has led some researchers to posit that intelligent biological organisms do not perceive time precisely and do not possess an explicit clock mechanism responsible for keeping track of it (Hills, 2003; Roseboom et al., 2019; Sherman et al., 2020).
Revised: Empirical evidence from the studies of human and animal cognition has led some researchers to posit that intelligent biological organisms do not perceive time precisely and do not possess an explicit clock mechanism responsible for keeping track of it (Hills, 2003; Roseboom et al., 2019; Sherman et al., 2020).

Original: 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 ﬁeld).
Revised: 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 ﬁeld).

Original: 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).
Revised: 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).

Original: Inspired by this account, in this work we propose subjectivetimescale model (STM), an alternative approach to learning a transition dynamics model, by replacing the objective timescale with a subjective one.
Revised: Inspired by this account, in this work we propose subjectivetimescale model (STM), an alternative approach to learning a transition dynamics model, by replacing the objective timescale with a subjective one.

Original: The latter represents the timescale by which an agent perceives events in an environment, predicts future states, and which is deﬁned by the sequences of episodic memories.
Revised: The latter represents the timescale by which an agent perceives events in an environment, predicts future states, and which is deﬁned by the sequences of episodic memories.

Original: 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 results in an agent with an ability to plan over varying timescales and construct novel future scenarios.
Revised: 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 results in an agent with an ability to plan over varying timescales and construct novel future scenarios.

Original: We employ two disparate agents to analyse the emergent characteristics of STM and demonstrate its applicability and usefulness to a variety of model-based agents.
Revised: We employ two disparate agents to analyse the emergent characteristics of STM and demonstrate its applicability and usefulness to a variety of model-based agents.

Original: We particularly emphasise the ease with which STM can be plugged into an existing model-based algorithm to signiﬁcantly improve its performance.
Revised: We particularly emphasise the ease with which STM can be plugged into an existing model-based algorithm to signiﬁcantly improve its performance.

Original: For this work, we use Dreamer by
Revised: For this work, we use Dreamer by

Original: Hafner et al.
Revised: Hafner et al.

Original: (2020) and deep active inference by Fountas et al.
Revised: (2020) and deep active inference by Fountas et al.

Original: (2020a).
Revised: (2020a).

Original: While these agents involve considerably different training and planning procedures, their underlying nature implies the central role of a generative model, including the learned transition dynamics model.
Revised: While these agents involve considerably different training and planning procedures, their underlying nature implies the central role of a generative model, including the learned transition dynamics model.

Original: Thus, as will be seen, both of the agents beneﬁt from the incorporation of
Revised: Thus, as will be seen, both of the agents beneﬁt from the incorporation of

Original: STM.
Revised: STM.

Original: Using these two agents and a combination of visuallycomplex ﬁrst-person foraging tasks, 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.
Revised: Using these two agents and a combination of visuallycomplex ﬁrst-person foraging tasks, 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.

Original: 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.
Revised: 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.

Original: Additionally, being trained on salient events, STM much more frequently imagines futures that contain epistemically-surprising events, which incentivises exploratory behaviour.
Revised: Additionally, being trained on salient events, STM much more frequently imagines futures that contain epistemically-surprising events, which incentivises exploratory behaviour.

Original: In both cases, the STM version outperforms the objective-timescale models, achieving higher rewards with signiﬁcantly fewer penalties.
Revised: In both cases, the STM version outperforms the objective-timescale models, achieving higher rewards with signiﬁcantly fewer penalties.

Original: 2. Related Work
Revised: 2. Related Work

Original: Model-based RL.
Revised: Model-based RL.

Original: Internal models are extensively studied in the ﬁeld of model-based RL.
Revised: Internal models are extensively studied in the ﬁeld of model-based RL.

Original: Using linear models to explicitly model transition dynamics has achieved impressive results in robotics (Bagnell & Schneider, 2001;
Revised: Using linear models to explicitly model transition dynamics has achieved impressive results in robotics (Bagnell & Schneider, 2001;

Original: Abbeel et al., 2006; Levine & Abbeel, 2014a;b; Watter et al., 2015; Levine et al., 2016; Kumar et al., 2016).
Revised: Abbeel et al., 2006; Levine & Abbeel, 2014a;b; Watter et al., 2015; Levine et al., 2016; Kumar et al., 2016).

Original: In general, however, their application is limited to low-dimensional domains and relatively simple environment dynamics.
Revised: In general, however, their application is limited to low-dimensional domains and relatively simple environment dynamics.

Original: Similarly, Gaussian Processes (GPs) have been used (Ko et al., 2007; Deisenroth & Rasmussen, 2011).
Revised: Similarly, Gaussian Processes (GPs) have been used (Ko et al., 2007; Deisenroth & Rasmussen, 2011).

Original: Their probabilistic nature allows for state uncertainty estimation, which can be incorporated in the planning module to make more cautious predictions; however, GPs struggle to scale to highdimensional data.
Revised: Their probabilistic nature allows for state uncertainty estimation, which can be incorporated in the planning module to make more cautious predictions; however, GPs struggle to scale to highdimensional data.

Original: An alternative and recently more prevalent method for parametrising transition models is to use neural networks.
Revised: An alternative and recently more prevalent method for parametrising transition models is to use neural networks.

Original: These are particularly attractive due to their recent proven success in a variety of domains, including deep model-free RL (Silver et al., 2017), ability to deal with high-dimensional data, and existence of methods for uncertainty quantiﬁcation (Blundell et al., 2015; Gal & Ghahramani, 2016).
Revised: These are particularly attractive due to their recent proven success in a variety of domains, including deep model-free RL (Silver et al., 2017), ability to deal with high-dimensional data, and existence of methods for uncertainty quantiﬁcation (Blundell et al., 2015; Gal & Ghahramani, 2016).

Original: Different deep learning architectures have been utilised including fully-connected neural networks (Nagabandi et al., 2018; Feinberg et al., 2018;
Revised: Different deep learning architectures have been utilised including fully-connected neural networks (Nagabandi et al., 2018; Feinberg et al., 2018;

Original: Kurutach et al., 2018) and autoregressive models (Racani ` ere et al., 2017; Ha & Schmidhuber, 2018; Ke et al., 2019), showing promising results in environments with relatively high-dimensional state spaces.
Revised: Kurutach et al., 2018) and autoregressive models (Racani ` ere et al., 2017; Ha & Schmidhuber, 2018; Ke et al., 2019), showing promising results in environments with relatively high-dimensional state spaces.

Original: In particular, autoregressive architectures, such as Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997), are capable of modelling non-Markovian environments and of learning temporal dependencies.
Revised: In particular, autoregressive architectures, such as Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997), are capable of modelling non-Markovian environments and of learning temporal dependencies.

Original: Nevertheless, LSTMs are still limited in their ability to learn relations between temporally-distant events, which is exacerbated in environments where little change occurs given an action.
Revised: Nevertheless, LSTMs are still limited in their ability to learn relations between temporally-distant events, which is exacerbated in environments where little change occurs given an action.

Original: Uncertainty quantiﬁcation using ensemble methods (Kalweit & Boedecker, 2017; Buckman et al., 2018; Clavera et al., 2020) or Bayesian neural networks (McAllister & Rasmussen, 2016; Depeweg et al., 2017) have been proposed to tackle model bias and sequential error accumulation.
Revised: Uncertainty quantiﬁcation using ensemble methods (Kalweit & Boedecker, 2017; Buckman et al., 2018; Clavera et al., 2020) or Bayesian neural networks (McAllister & Rasmussen, 2016; Depeweg et al., 2017) have been proposed to tackle model bias and sequential error accumulation.

Original: Other works have focused on techniques to create more accurate long-term predictions.
Revised: Other works have focused on techniques to create more accurate long-term predictions.

Original: Mishra et al.
Revised: Mishra et al.

Original: (2017) used a segmentbased approach to predict entire trajectories at once in an attempt to avoid one-step prediction error accumulation.
Revised: (2017) used a segmentbased approach to predict entire trajectories at once in an attempt to avoid one-step prediction error accumulation.

Original: A work by Ke et al.
Revised: A work by Ke et al.

Original: (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-horizon predictions.
Revised: (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-horizon predictions.

Original: In contrast, the work presented in this paper re-focuses the objective from attempting to create better parametrisation techniques or mitigating methods to simply transforming the timescale over which the dynamics of an environment is learned.
Revised: In contrast, the work presented in this paper re-focuses the objective from attempting to create better parametrisation techniques or mitigating methods to simply transforming the timescale over which the dynamics of an environment is learned.

Original: As will be seen, our approach can lead to more informative and efﬁcient long-term predictions without compromising agent’s ability to plan over short time-horizons.
Revised: As will be seen, our approach can lead to more informative and efﬁcient long-term predictions without compromising agent’s ability to plan over short time-horizons.

Original: Time and Models.
Revised: Time and Models.

Original: Relatively little work has been done in the ﬁeld to study the role and practical use of time perception for artiﬁcial agents, despite its apparent importance for intelligent behaviour (Deverett et al., 2019).
Revised: Relatively little work has been done in the ﬁeld to study the role and practical use of time perception for artiﬁcial agents, despite its apparent importance for intelligent behaviour (Deverett et al., 2019).

Original: For instance, Braylan et al.
Revised: For instance, Braylan et al.

Original: (2015) showed that the degree of frame-skipping can have a signiﬁcant impact on an agent’s performance in Atari 2600 games.
Revised: (2015) showed that the degree of frame-skipping can have a signiﬁcant impact on an agent’s performance in Atari 2600 games.

Original: Nevertheless, several works have proposed methods for learning models with particular temporal characteristics, such as time-agnostic models (Jayaraman et al., 2019) or ‘jumpy’ models with ﬁxed (Gregor & Besse, 2019; Buesing et al., 2018) or adaptive frame intervals (Goyal et al., 2019; Neitz et al., 2018;
Revised: Nevertheless, several works have proposed methods for learning models with particular temporal characteristics, such as time-agnostic models (Jayaraman et al., 2019) or ‘jumpy’ models with ﬁxed (Gregor & Besse, 2019; Buesing et al., 2018) or adaptive frame intervals (Goyal et al., 2019; Neitz et al., 2018;

Original: Pertsch et al., 2020).
Revised: Pertsch et al., 2020).

Original: In line with some of these works, our approach is designed to perform adaptive frame skips, but unlike all of them it is characterised by the simplicity of use and incorporation into existing model-based methods.
Revised: In line with some of these works, our approach is designed to perform adaptive frame skips, but unlike all of them it is characterised by the simplicity of use and incorporation into existing model-based methods.

Original: Furthermore, our method aims to mimic human time perception, being inspired by the latest works from computational neuroscience.
Revised: Furthermore, our method aims to mimic human time perception, being inspired by the latest works from computational neuroscience.

Original: Active Inference.
Revised: Active Inference.

Original: Until now, most of the work on active inference has been done in low-dimensional and discrete state spaces (Friston et al., 2015; 2017a;b;c).
Revised: Until now, most of the work on active inference has been done in low-dimensional and discrete state spaces (Friston et al., 2015; 2017a;b;c).

Original: Recently, however, there has been a rising interest in scaling active inference and applying it to environments with continuous and/or large state spaces (Ueltzh ¨ offer, 2018; Tschantz et al., 2019; C¸atal et al., 2019; Millidge, 2019; Fountas et al., 2020a).
Revised: Recently, however, there has been a rising interest in scaling active inference and applying it to environments with continuous and/or large state spaces (Ueltzh¨offer, 2018; Tschantz et al., 2019; C¸atal et al., 2019; Millidge, 2019; Fountas et al., 2020a).

Original: 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.
Revised: 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.

Original: Episodic Memory.
Revised: Episodic Memory.

Original: In neuroscience, episodic memory is used to describe autobiographical memories that link a collection of ﬁrst-person sensory experiences at a speciﬁc time and place (Tulving, 1972).
Revised: In neuroscience, episodic memory is used to describe autobiographical memories that link a collection of ﬁrst-person sensory experiences at a speciﬁc time and place (Tulving, 1972).

Original: Past studies in the ﬁeld suggest that episodic memory plays an important role in human learning (Mahr & Csibra, 2017), and may capture a wide range of potential functional purposes, such as construction of novel future scenarios (Schacter et al., 2007; Hassabis et al., 2007; Schacter et al., 2012), mental time-travel (Michaelian, 2016) or assisting in the formation of new semantic memories (Greenberg & Verfaellie, 2010).
Revised: Past studies in the ﬁeld suggest that episodic memory plays an important role in human learning (Mahr & Csibra, 2017), and may capture a wide range of potential functional purposes, such as construction of novel future scenarios (Schacter et al., 2007; Hassabis et al., 2007; Schacter et al., 2012), mental time-travel (Michaelian, 2016) or assisting in the formation of new semantic memories (Greenberg & Verfaellie, 2010).

Original: A recent computational model of episodic memory (Fountas et al., 2020b) also relates it to the human ability to estimate time durations.
Revised: A recent computational model of episodic memory (Fountas et al., 2020b) also relates it to the human ability to estimate time durations.

Original: The application of episodic memory in reinforcement learning has been somewhat limited.
Revised: The application of episodic memory in reinforcement learning has been somewhat limited.

Original: Some works have employed simple forms of memory to improve the performance of a deep model-free RL agent via experience replay (Mnih et al., 2015; Schaul et al., 2016; Espeholt et al., 2018).
Revised: Some works have employed simple forms of memory to improve the performance of a deep model-free RL agent via experience replay (Mnih et al., 2015; Schaul et al., 2016; Espeholt et al., 2018).

Original: However, these methods do not incorporate information about associative or temporal dependencies between the memories (Hansen et al., 2018).
Revised: However, these methods do not incorporate information about associative or temporal dependencies between the memories (Hansen et al., 2018).

Original: Read-write memory banks have also been implemented alongside gradient-based systems (memory-augmented neural networks) for assisting in learning and prediction (Graves et al., 2014; 2016;
Revised: Read-write memory banks have also been implemented alongside gradient-based systems (memory-augmented neural networks) for assisting in learning and prediction (Graves et al., 2014; 2016;

Original: Oh et al., 2016;
Revised: Oh et al., 2016;

Original: Jung et al., 2018; Hung et al., 2019).
Revised: Jung et al., 2018; Hung et al., 2019).

Original: Further, episodic memory has been used for non-parametric Q-function approximation (Hansen et al., 2018; Blundell et al., 2016; Pritzel et al., 2017; Zhu et al., 2020).
Revised: Further, episodic memory has been used for non-parametric Q-function approximation (Hansen et al., 2018; Blundell et al., 2016; Pritzel et al., 2017; Zhu et al., 2020).

Original: It has also been proposed to be used directly for control as a faster and more efﬁcient alternative to model-based and model-free approaches in RL, such as instance-based control (Lengyel & Dayan, 2007; Gershman & Daw, 2017; Botvinick et al., 2019) and one-shot learning (Kaiser et al., 2017).
Revised: It has also been proposed to be used directly for control as a faster and more efﬁcient alternative to model-based and model-free approaches in RL, such as instance-based control (Lengyel & Dayan, 2007; Gershman & Daw, 2017; Botvinick et al., 2019) and one-shot learning (Kaiser et al., 2017).

Original: In contrast, our paper considers a novel way of using episodic memories – in deﬁning the agent’s subjective timescale of the environment and training a transition dynamics model over the sequences of these memories.
Revised: In contrast, our paper considers a novel way of using episodic memories – in deﬁning the agent’s subjective timescale of the environment and training a transition dynamics model over the sequences of these memories.

Original: 3. Baseline Agents
Revised: 3. Baseline Agents

Original: To test the versatility of STM, we integrate it into two types of model-based agents that represent fundamentally different approaches to building model-based agents.
Revised: To test the versatility of STM, we integrate it into two types of model-based agents that represent fundamentally different approaches to building model-based agents.

Original: We consider a partially-observable Markov decision process (POMDP) with discrete time t .
Revised: We consider a partially-observable Markov decision process (POMDP) with discrete time t .

Original: At every time step, agents receive observation o t ∈
Revised: At every time step, agents receive observation o t ∈

Original: O
Revised: O

Original: upon taking a discrete action a t − 1 ∈ A .
Revised: upon taking a discrete action a t − 1 ∈ A .

Original: Latent states inferred by an agent are denoted by s t ∈ S .
Revised: Latent states inferred by an agent are denoted by s t ∈ S .

Original: For Dreamer, an agent also observes a scalar reward r t ∈ R .
Revised: For Dreamer, an agent also observes a scalar reward r t ∈ R .

Original: 3.1. Dreamer
Revised: 3.1. Dreamer

Original: Dreamer (Hafner et al., 2020) is a well-established modelbased RL agent that uses a recurrent generative model (Hafner et al., 2019) and latent imagination to learn a rewardmaximising policy.
Revised: Dreamer (Hafner et al., 2020) is a well-established modelbased RL agent that uses a recurrent generative model (Hafner et al., 2019) and latent imagination to learn a rewardmaximising policy.

Original: Speciﬁcally, it consists of: a representation model q ( s t | s t − 1 , a t − 1 , o t ; θ ) , a transition model p ( s t | s t − 1 , a t − 1 ; θ ) , a reward model q ( r t | s t ; θ ) , an action model q ( a t
Revised: Speciﬁcally, it consists of: a representation model q ( s t | s t − 1 , a t − 1 , o t ; θ ) , a transition model p ( s t | s t − 1 , a t − 1 ; θ ) , a reward model q ( r t | s t ; θ ) , an action model q ( a t

Original: | s t ; φ ) , and a value model v ( s t ; ψ ) (Fig. 2b), where { θ, φ, ψ } are parameters of the models.
Revised: | s t ; φ ) , and a value model v ( s t ; ψ ) (Fig. 2b), where { θ, φ, ψ } are parameters of the models.

Original: The dynamics model trained from environment observations is used to perform latent roll-outs up to horizon H while predicting rewards ˆ r t and values v ( s t ) , and then updating the action
Revised: The dynam-

Original: [Figure]
Revised: [Figure]

Original:  and value models using the value estimates V λ ( s t ) , such that the action model maximises expected rewards, while the value model regresses the value of a state:
Revised: ics model trained from environment observations is used to perform latent roll-outs up to horizon H while predicting rewards ˆ r t and values v ( s t ) , and then updating the action and value models using the value estimates V λ ( s t ) , such that the action model maximises expected rewards, while the value model regresses the value of a state:

Original: [Equation]
Revised: [Equation]

Original: In turn, value estimates V λ are calculated via a combination of predicted rewards and values (see Appendix A.2).
Revised: In turn, value estimates V λ are calculated via a combination of predicted rewards and values (see Appendix A.2).

Original: Dreamer performs competitively on a wide range of traditional RL tasks.
Revised: Dreamer performs competitively on a wide range of traditional RL tasks.

Original: 3.2. Deep Active Inference
Revised: 3.2. Deep Active Inference

Original: The second baseline is a deep active inference (DAI) agent with Monte-Carlo tree search (Fountas et al., 2020a).
Revised: The second baseline is a deep active inference (DAI) agent with Monte-Carlo tree search (Fountas et al., 2020a).

Original: Its selection is motivated by an observation-based objective and a purely model-based planning procedure, which allow for a more detailed analysis of STM’s properties.
Revised: Its selection is motivated by an observation-based objective and a purely model-based planning procedure, which allow for a more detailed analysis of STM’s properties.

Original: The generative model of this agent is deﬁned as p ( o 1: t , s 1: t , a 1: t ; θ ) , where θ are parameters of the model.
Revised: The generative model of this agent is deﬁned as p ( o 1: t , s 1: t , a 1: t ; θ ) , where θ are parameters of the model.

Original: Speciﬁcally, it includes two factors: a transition model p ( s t | s t − 1 , a t − 1 ; θ ) and a latent state decoder p ( o t | s t ; θ ) parametrised by feed-forward neural networks (Fig. 2a).
Revised: Speciﬁcally, it includes two factors: a transition model p ( s t | s t − 1 , a t − 1 ; θ ) and a latent state decoder p ( o t | s t ; θ ) parametrised by feed-forward neural networks (Fig. 2a).

Original: The agent also possesses two inference models, which are trained using amortized inference: a habitual network q ( a t ; φ a ) and observation encoder q ( s t ; φ s ) parametrised by φ a and φ s , respectively.
Revised: The agent also possesses two inference models, which are trained using amortized inference: a habitual network q ( a t ; φ a ) and observation encoder q ( s t ; φ s ) parametrised by φ a and φ s , respectively.

Original: Further, we use π to denote a sequence of actions up to some time horizon H .
Revised: Further, we use π to denote a sequence of actions up to some time horizon H .

Original: Training is performed using the variational free energy at an arbitrary time-step t :
Revised: Training is performed using the variational free energy at an arbitrary time-step t :

Original: [Equation]
Revised: [Equation]

Original: where p ( a ) =
Revised: where p ( a ) =

Original: (cid:80) π : a 1 = a p ( π ) is the summed probability of all policies π beginning with action a .
Revised: (cid:80) π : a 1 = a p ( π ) is the summed probability of all policies π beginning with action a .

Original: The expected free energy (EFE) of the generative model up to some time horizon H is deﬁned as:
Revised: The expected free energy (EFE) of the generative model up to some time horizon H is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: where ˜ q 
Revised: where ˜ q = q

Original: [Equation] π ) and p ( o k , s k , θ | π ) =
Revised: ( o k , s k , θ | π ) and p ( o k , s k , θ | π ) =

Original: p ( o k | π ) q ( s k |
Revised: p ( o k | π ) q ( s k |

Original: o k , π ) p ( θ | s k , o k , π ) .
Revised: o k , π ) p ( θ | s k , o k , π ) .

Original: Action selection is aided with Monte Carlo tree search (MCTS), ensuring a more efﬁcient trajectory search.
Revised: Action selection is aided with Monte Carlo tree search (MCTS), ensuring a more efﬁcient trajectory search.

Original: Speciﬁcally, MCTS generates a weighted tree that is used to sample policies from the current timestep, where the weights refer to the agent’s estimation of the EFE given a state-action pair, ˜ G ( s, a ) .
Revised: Speciﬁcally, MCTS generates a weighted tree that is used to sample policies from the current timestep, where the weights refer to the agent’s estimation of the EFE given a state-action pair, ˜ G ( s, a ) .

Original: The nodes of the tree are predicted via the transition model, p ( s t | s t − 1 , a t − 1 ; θ ) .
Revised: The nodes of the tree are predicted via the transition model, p ( s t | s t − 1 , a t − 1 ; θ ) .

Original: At the end of the search,
Revised: At the end of the search,

Original: MCTS is used to construct the action prior, p ( a t ) =
Revised: MCTS is used to construct the action prior, p ( a t ) =

Original: N ( a t , s t ) /
Revised: N ( a t , s t ) /

Original: (cid:80) j N ( a j,t , s t ) , where N ( a t , s t ) is the number of times action a t has been taken from state s t .
Revised: (cid:80) j N ( a j,t , s t ) , where N ( a t , s t ) is the number of times action a t has been taken from state s t .

Original: More details on active inference and MCTS planning procedure can be found in Appendix A and C.2.
Revised: More details on active inference and MCTS planning procedure can be found in Appendix A and C.2.

Original: 4. Subjective-Timescale Model
Revised: 4. Subjective-Timescale Model

Original: We introduce subjective-timescale model (STM) – a transition dynamics model trained over sequences of episodic memories that deﬁne a more useful timescale.
Revised: We introduce subjective-timescale model (STM) – a transition dynamics model trained over sequences of episodic memories that deﬁne a more useful timescale.

Original: As such, the system consists of a memory accumulation module to selectively record salient events and an autoregressive transition model (Figure 3).
Revised: As such, the system consists of a memory accumulation module to selectively record salient events and an autoregressive transition model (Figure 3).

Original: We deﬁne a ground-truth sequence as a sequence of all states experienced in an environment during a single episode denoted with t , S g
Revised: We deﬁne a ground-truth sequence as a sequence of all states experienced in an environment during a single episode denoted with t , S g

Original: = { s 0 , ..., s t , s t +1 , ..., s T } , and an S-sequence (subjective sequence) as a sequence of states selectively picked by our system, and over which the new transition model would be learned,
Revised: = { s 0 , ..., s t , s t +1 , ..., s T } , and an S-sequence (subjective sequence) as a sequence of states selectively picked by our system, and over which the new transition model would be learned,

Original: S e = { s 0 , ..., s τ , s τ +1 , ..., s T } , denoted with τ .
Revised: S e = { s 0 , ..., s τ , s τ +1 , ..., s T } , denoted with τ .

Original: Each unit in an Ssequence is called an episodic memory and consists of a set of sufﬁcient statistics, s = { µ s , σ s } , where µ s and σ s are mean and variance vectors of a Gaussian-distributed state s , respectively.
Revised: Each unit in an Ssequence is called an episodic memory and consists of a set of sufﬁcient statistics, s = { µ s , σ s } , where µ s and σ s are mean and variance vectors of a Gaussian-distributed state s , respectively.

Original: Additionally, each episodic memory contains a reference to its preceding (parent) episodic memory and all actions until the next one.
Revised: Additionally, each episodic memory contains a reference to its preceding (parent) episodic memory and all actions until the next one.

Original: We call the process of recording
Revised: We call the process of recording

Original: S-sequences memory accumulation .
Revised: S-sequences memory accumulation .

Original: 4.1. Memory Accumulation
Revised: 4.1. Memory Accumulation

Original: Previous work on time perception and episodic memory (Fountas et al., 2020b) employed saliency of an event, or the generative model’s prediction error, as the memory formation criterion.
Revised: Previous work on time perception and episodic memory (Fountas et al., 2020b) employed saliency of an event, or the generative model’s prediction error, as the memory formation criterion.

Original: Selection of this criterion is informed by the experimental evidence from neuroscience on episodic memory (Greve et al., 2017; Jang et al., 2018;
Revised: Selection of this criterion is informed by the experimental evidence from neuroscience on episodic memory (Greve et al., 2017; Jang et al., 2018;

Original: Rouhani et al., 2018) and recent time perception models (Roseboom et al., 2019).
Revised: Rouhani et al., 2018) and recent time perception models (Roseboom et al., 2019).

Original: Inspired by this account, our memory accumulation system uses the prediction error (or the free energy) of the objective-timescale transition model as a measure of event saliency, and forms memories when a pre-deﬁned threshold is exceeded.
Revised: Inspired by this account, our memory accumulation system uses the prediction error (or the free energy) of the objective-timescale transition model as a measure of event saliency, and forms memories when a pre-deﬁned threshold is exceeded.

Original: Accumulation To train STM, an agent moves in the environment under a pre-trained generative model.
Revised: Accumulation To train STM, an agent moves in the environment under a pre-trained generative model.

Original: During this process, each transition is evaluated based on the objective transition model prediction error (or free energy), which represents the degree of surprise experienced by the transition model upon taking an action.
Revised: During this process, each transition is evaluated based on the objective transition model prediction error (or free energy), which represents the degree of surprise experienced by the transition model upon taking an action.

Original: If the value exceeds a predeﬁned threshold, (cid:15) , a memory is formed and placed into an S-sequence.
Revised: If the value exceeds a predeﬁned threshold, (cid:15) , a memory is formed and placed into an S-sequence.

Original: At the end of each episode, the recorded Ssequence is saved for later use.
Revised: At the end of each episode, the recorded Ssequence is saved for later use.

Original: Over the course of training, we increase the threshold using (cid:15) =
Revised: Over the course of training, we increase the threshold using (cid:15) =

Original: (cid:15) max (1 − exp( − λ · i )) , where (cid:15) max is the maximum allowable threshold, λ is the half-life of the exponential decay term and i is training iteration.
Revised: (cid:15) max (1 − exp( − λ · i )) , where (cid:15) max is the maximum allowable threshold, λ is the half-life of the exponential decay term and i is training iteration.

Original: Recorded memories We can categorise the transitions that cause higher values of transition model prediction errors into two main groups: epistemic surprise and modelimperfection surprise.
Revised: Recorded memories We can categorise the transitions that cause higher values of transition model prediction errors into two main groups: epistemic surprise and modelimperfection surprise.

Original: 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).
Revised: 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).

Original: The latter refers to the main bulk of these high prediction-error transitions and stems from the inherent imperfections of the learned dynamics.
Revised: The latter refers to the main bulk of these high prediction-error transitions and stems from the inherent imperfections of the learned dynamics.

Original: Speciﬁcally, less frequently-occurring observations with richer combinatorial structure would systematically result in higher compounded transition model errors, given that these would be characteristic of more complex scenes.
Revised: Speciﬁcally, less frequently-occurring observations with richer combinatorial structure would systematically result in higher compounded transition model errors, given that these would be characteristic of more complex scenes.

Original: Empirical evidence of these trends can be found in Appendix D. As will become apparent, the presence of these two categories in the recorded S-sequences results in the model’s ability to vary its prediction timescale based on the perceptual context and systematically imagine future salient events.
Revised: Empirical evidence of these trends can be found in Appendix D. As will become apparent, the presence of these two categories in the recorded S-sequences results in the model’s ability to vary its prediction timescale based on the perceptual context and systematically imagine future salient events.

Original: Two examples of S-sequences are shown in Figure 4, where colours are used to indicate parts of the sequences with different accumulation rates.
Revised: Two examples of S-sequences are shown in Figure 4, where colours are used to indicate parts of the sequences with different accumulation rates.

Original: A transition dynamics model is necessarily trained with respect to actions that an agent took to reach subsequent states.
Revised: A transition dynamics model is necessarily trained with respect to actions that an agent took to reach subsequent states.

Original: However, STM records memories over an arbitrary number of steps, thus leaving action sequences of variable length.
Revised: However, STM records memories over an arbitrary number of steps, thus leaving action sequences of variable length.

Original: For the purposes of this paper, we implement a simple heuristic to summarise agent’s trajectories, which is enough to provide our system with the necessary information to learn action-conditioned predictions.
Revised: For the purposes of this paper, we implement a simple heuristic to summarise agent’s trajectories, which is enough to provide our system with the necessary information to learn action-conditioned predictions.

Original: We do it by estimating the angle between the agent’s initial position and its ﬁnal position at the time-step of the subsequent memory (full details can be found in Appendix B.3).
Revised: We do it by estimating the angle between the agent’s initial position and its ﬁnal position at the time-step of the subsequent memory (full details can be found in Appendix B.3).

Original: 4.2. Transition Dynamics Model
Revised: 4.2. Transition Dynamics Model

Original: As mentioned, S-sequences are characterised by the presence of epistemically-surprising and salient events squeezed together in the recorded episodes.
Revised: As mentioned, S-sequences are characterised by the presence of epistemically-surprising and salient events squeezed together in the recorded episodes.

Original: As a result, training on these sequences is more conducive for learning temporal dependencies between important states.
Revised: As a result, training on these sequences is more conducive for learning temporal dependencies between important states.

Original: For this reason, we employ a recurrent transition dynamics model (LSTM (Hochreiter & Schmidhuber, 1997) for active inference and
Revised: For this reason, we employ a recurrent transition dynamics model (LSTM (Hochreiter & Schmidhuber, 1997) for active inference and

Original: GRU (Cho et al., 2014) for Dreamer), which utilises internal memory states to store information about preceding inputs.
Revised: GRU (Cho et al., 2014) for Dreamer), which utilises internal memory states to store information about preceding inputs.

Original: As such, hidden states h τ at subjective time τ are calculated via a deterministic mapping, h τ = f θ h
Revised: As such, hidden states h τ at subjective time τ are calculated via a deterministic mapping, h τ = f θ h

Original: ( s τ , a τ , h τ − 1 ) , where s τ and a τ are the latent state and action, respectively, and f θ h is a deterministic function of the used recurrent unit with parameters θ h .
Revised: ( s τ , a τ , h τ − 1 ) , where s τ and a τ are the latent state and action, respectively, and f θ h is a deterministic function of the used recurrent unit with parameters θ h .

Original: The hidden state h τ is then mapped to a latent state s τ +1 at the next subjective time τ + 1 .
Revised: The hidden state h τ is then mapped to a latent state s τ +1 at the next subjective time τ + 1 .

Original: Importantly, the parameters of STM can be trained normally according to the procedure of the model-based agent or in isolation via loss function deﬁned as,
Revised: Importantly, the parameters of STM can be trained normally according to the procedure of the model-based agent or in isolation via loss function deﬁned as,

Original: [Equation]
Revised: [Equation]

Original: The STM generative model is shown in Figure 2c and the procedure is outlined in Algorithm 1.
Revised: The STM generative model is shown in Figure 2c and the procedure is outlined in Algorithm 1.

Original: Architectural and training details of the model can be found in Appendix B.
Revised: Architectural and training details of the model can be found in Appendix B.

Original: 5. Experiments
Revised: 5. Experiments

Original: To analyse the properties and beneﬁts of STM, we use a combination of visually-complex 3D foraging tasks from
Revised: To analyse the properties and beneﬁts of STM, we use a combination of visually-complex 3D foraging tasks from

Original: DeepMind Lab (DMLab) (Beattie et al., 2016) and AnimalAI (AAI) (Crosby et al., 2020; Crosby, 2020).
Revised: DeepMind Lab (DMLab) (Beattie et al., 2016) and AnimalAI (AAI) (Crosby et al., 2020; Crosby, 2020).

Original: DMLab and AAI afford suitable testing conditions given their large size, slow objective temporal dynamics, and low density of rewards.
Revised: DMLab and AAI afford suitable testing conditions given their large size, slow objective temporal dynamics, and low density of rewards.

Original: Foraging with sparse rewards Within DMLab and AAI, we use sparsely populated arenas with positive 2 and negative rewards.
Revised: Foraging with sparse rewards Within DMLab and AAI, we use sparsely populated arenas with positive 2 and negative rewards.

Original: The chosen map conﬁguration is motivated by the Algorithm
Revised: The chosen map conﬁguration is motivated by the fact that a successful agent would be forced to perform both short- and long-distance planning, as well as a more extensive exploration of the environment.

Original: Initialise :
Revised: Foraging in a maze Further, we test the Dreamer agent within a more complex maze conﬁguration in DMLab. These maps include additional levels of complexity such as epistemic distractions (wall imagery) and more difﬁcult

Original: Pre-trained objective-timescale generative model, p ( o, s, a ; θ )
Revised: Pre-trained objective-timescale generative model, p ( o, s, a ; θ )

Original: Memory threshold, (cid:15)
Revised: Memory threshold, (cid:15)

Original: Episodic buffer, D ← {} for episode = 1 to ∞ do
Revised: Episodic buffer, D ← {} for episode = 1 to ∞ do

Original: Initialise S-sequence with the ﬁrst observation and state,
Revised: Initialise S-sequence with the ﬁrst observation and state,

Original: S ← { s 1 , o 1 } for t = 1 to max length do //
Revised: S ← { s 1 , o 1 } for t = 1 to max length do //

Original: Interact and collect memories
Revised: Interact and collect memories

Original: Sample an action and get a new observation
Revised: Sample an action and get a new observation

Original: Compute surprise of the objective-timescale transition model, L obj
Revised: Compute surprise of the objective-timescale transition model, L obj

Original: if L obj ≥ (cid:15) then
Revised: if L obj ≥ (cid:15) then

Original: Add memory: S ← S ∪ { a t , s t +1 , o t +1 } else
Revised: Add memory: S ← S ∪ { a t , s t +1 , o t +1 } else

Original: Add action: S ← S ∪ { a t } // Training
Revised: Add action: S ← S ∪ { a t } // Training

Original: Sample a batch of S-sequences: B i ∼ D
Revised: Sample a batch of S-sequences: B i ∼ D

Original: Train the models: E B i
Revised: Train the models: E B i

Original: [ L ]fact that a successful agent would be forced to perform both short- and long-distance planning, as well as a more extensive exploration of the environment. Foraging in a maze Further, we test the Dreamer agent within a more complex maze conﬁguration in DMLab. These maps include additional levels of complexity such as epistemic distractions (wall imagery) and more difﬁcult terrain with obstructions.
Revised: [ L ] terrain with obstructions.

Original: More pronounced exploratory navigation and reward signals are crucial in this setting.
Revised: More pronounced exploratory navigation and reward signals are crucial in this setting.

Original: Training Across all tasks, we use the environment’s intrinsic frame rate (i.e. no action repetitions) and an episode length of 1000 steps.
Revised: Training Across all tasks, we use the environment’s intrinsic frame rate (i.e. no action repetitions) and an episode length of 1000 steps.

Original: Observations are of size 32 × 32 ×and 64 × 64 × 3 for AAI and DMLab, respectively.
Revised: Observations are of size 32 × 32 ×and 64 × 64 × 3 for AAI and DMLab, respectively.

Original: We train the objective- and subjective-timescale agents using the same hyperparameters on a single NVIDIA Quadro RTX 6000 GPU.[Table] [Table] [Table]
Revised: We train the objective- and subjective-timescale agents using the same hyperparameters on a single NVIDIA Quadro RTX 6000 GPU.

Original: [Table]

Original: 5.1. Reward-based Results
Revised: 5.1. Reward-based Results

Original: Enhanced reward seeking Upon training completion
Revised: Enhanced reward seeking Upon training completion

Original: , we compare the objective- and subjective-timescale agents performance in randomly-generated environments and validate a statistically signiﬁcant improvement in solving the tasks by agents with STM.
Revised: , we compare the objective- and subjective-timescale agents performance in randomly-generated environments and validate a statistically signiﬁcant improvement in solving the tasks by agents with STM.

Original: Table 1 summarises reward-based results.
Revised: Table 1 summarises reward-based results.

Original: This empirical evidence indicates that training and planning over the subjective timescale can be considerably advantageous to model-based agents.
Revised: This empirical evidence indicates that training and planning over the subjective timescale can be considerably advantageous to model-based agents.

Original: More cautious behaviour The ability of STM to perform both short- and long-term planning allows the agent to be more cautious in selecting actions that may lead to negative rewards in the future.
Revised: More cautious behaviour The ability of STM to perform both short- and long-term planning allows the agent to be more cautious in selecting actions that may lead to negative rewards in the future.

Original: Table 1 shows a large drop in the collected negative rewards when the agents employ STM.
Revised: Table 1 shows a large drop in the collected negative rewards when the agents employ STM.

Original: 5.2. Varying Prediction Timescale
Revised: 5.2. Varying Prediction Timescale

Original: Much like human perception of time changes depending on the perceptual content of the surroundings, our agent  varies the prediction timescale depending on the context it ﬁnds itself in.
Revised: Much like human perception of time changes depending on the perceptual content of the surroundings, our agent [Table] varies the prediction timescale depending on the context it ﬁnds itself in.

Original: Speciﬁcally, in a static environment the complexity of any given observation is primarily driven by its conﬁgurable parts (e.g. objects, which may appear in different sizes, colours, locations, etc.).
Revised: Speciﬁcally, in a static environment the complexity of any given observation is primarily driven by its conﬁgurable parts (e.g. objects, which may appear in different sizes, colours, locations, etc.).

Original: As a result, our agent tends to predict farther into the future in the absence of any nearby objects, and slows its timescale, predicting at ﬁner temporal rate, when the objects are close.
Revised: As a result, our agent tends to predict farther into the future in the absence of any nearby objects [Table] and slows its timescale, predicting at ﬁner temporal rate, when the objects are close.

Original: The EFE planning objective of the active inference agent allows us to visualise this property of STM by means of an energy map.
Revised: The EFE planning objective of the active inference agent allows us to visualise this property of STM by means of an energy map.

Original: Figure 6a shows the agent’s values of EFE at different locations in the environment calculated using the active inference agent with and without STM using identical MCTS parameters.
Revised: Figure 6a shows the agent’s values of EFE at different locations in the environment calculated using the active inference agent with and without STM using identical MCTS parameters.

Original: Speciﬁcally, the map contains a single green sphere placed in the middle, and the agent’s orientation is always towards the top of the image.
Revised: Speciﬁcally, the map contains a single green sphere placed in the middle, and the agent’s orientation is always towards the top of the image.

Original: Regions with lower values of EFE (more blue) indicate the agent’s ability to ‘see’ the sphere in its imagination roll-outs during planning, since its appearance decreases the value of the observational prior component.
Revised: Regions with lower values of EFE (more blue) indicate the agent’s ability to ‘see’ the sphere in its imagination roll-outs during planning, since its appearance decreases the value of the observational prior component.

Original: The comparison of these regions produced by the two models indicate two important things.
Revised: The comparison of these regions produced by the two models indicate two important things.

Original: First, the signiﬁcantly larger blue region of DAI with STM (right) points to the model’s ability to predict farther into the future, when placed at a distance from the object of interest, in contrast to the agent with the objectivetimescale transition model.
Revised: First, the signiﬁcantly larger blue region of DAI with STM (right) points to the model’s ability to predict farther into the future, when placed at a distance from the object of interest, in contrast to the agent with the objectivetimescale transition model.

Original: Second, STM’s blue region also extends to the regions very close to the sphere, suggesting that short-term planning is not compromised.
Revised: Second, STM’s blue region also extends to the regions very close to the sphere, suggesting that short-term planning is not compromised.

Original: Performing temporal jumps and skipping unnecessary intermediary steps affords greater computational efﬁciency and reduces the detrimental effects of sequential error accumulation, as can be seen in Figure 5.
Revised: Performing temporal jumps and skipping unnecessary intermediary steps affords greater computational efﬁciency and reduces the detrimental effects of sequential error accumulation, as can be seen in Figure 5.

Original: Similarly, while
Revised: Similarly, while

Original: STM is able to predict far ahead, its inherent ﬂexibility to automatically predict over varying timescales does not compromise the agent’s performance when the states of interest
Revised: STM is able to predict far ahead, its inherent ﬂexibility to automatically predict over varying timescales does not com-

Original: [Figure]
Revised: [Figure]

Original:  are close.
Revised: promise the agent’s performance when the states of interest are close.

Original: Thus, a separate mechanism for adjusting how far into the future an agent should plan is not necessary and is implicitly handled by our model.
Revised: Thus, a separate mechanism for adjusting how far into the future an agent should plan is not necessary and is implicitly handled by our model.

Original: Finally, STM allows the agent to make more informed decisions in an environment, as it tends to populate the roll-outs with salient observations of the short- and long-term futures depending on the context.
Revised: Finally, STM allows the agent to make more informed decisions in an environment, as it tends to populate the roll-outs with salient observations of the short- and long-term futures depending on the context.

Original: As such, STM effectively re-focuses the central purpose of a transition model from most accurately modeling the ground-truth dynamics of an environment to predicting states more informative with respect to the affordances of its surroundings.
Revised: As such, STM effectively re-focuses the central purpose of a transition model from most accurately modeling the ground-truth dynamics of an environment to predicting states more informative with respect to the affordances of its surroundings.

Original: 5.3. Imagining Surprising Events
Revised: 5.3. Imagining Surprising Events

Original: In the absence of interesting surroundings, it is useful for an agent to imagine the kinds of observations it could see in the future.
Revised: In the absence of interesting surroundings, it is useful for an agent to imagine the kinds of observations it could see in the future.

Original: We show that since S-sequences frequently include epistemically-surprising memories, STM favours prediction of salient events.
Revised: We show that since S-sequences frequently include epistemically-surprising memories, STM favours prediction of salient events.

Original: For instance, in the context of the selected environments, these events mostly constitute objects coming into agent’s view upon turning.
Revised: For instance, in the context of the selected environments, these events mostly constitute objects coming into agent’s view upon turning.

Original: To investigate this, we compared 1000 roll-outs using DAI with and without STM (example shown in Fig.7), counting objects that the models imagine starting from an initial observation of an empty arena.
Revised: To investigate this, we compared 1000 roll-outs using DAI with and without STM (example shown in Fig.7), counting objects that the models imagine starting from an initial observation of an empty arena.

Original: For roll-outs conditioned on turning right, we detect 1274 objects in roll-outs made with STM compared towithout STM; similarly, for turning left, 941 with STM compared to 0 without STM.
Revised: For roll-outs conditioned on turning right, we detect 1274 objects in roll-outs made with STM compared towithout STM; similarly, for turning left, 941 with STM compared to 0 without STM.

Original: The ability of the STM to imagine novel and salient future events encourages exploratory behaviour that is distinct from the agents’ information seeking motivations.
Revised: The ability of the STM to imagine novel and salient future events encourages exploratory behaviour that is distinct from the agents’ information seeking motivations.

Original: To test this, we create an empty AAI arena with light brown walls and set a pre-trained active inference agent’s observational prior to a green colour.
Revised: To test this, we create an empty AAI arena with light brown walls and set a pre-trained active inference agent’s observational prior to a green colour.

Original: In this scenario, brown walls act as suboptimal attractors that reduce agent’s EFE, meaning that the agent is likely to stick around them unless it can imagine more optimal states elsewhere.
Revised: In this scenario, brown walls act as suboptimal attractors that reduce agent’s EFE, meaning that the agent is likely to stick around them unless it can imagine more optimal states elsewhere.

Original: Figure 6b shows a striking difference in the heat map of visited locations between DAI without and with STM (after 50k steps).
Revised: Figure 6b shows a striking difference in the heat map of visited locations between DAI without and with STM (after 50k steps).

Original: As can be seen, the
Revised: As can be seen, the

Original: STM agent explores the environment even in the presence of the ubiquitous suboptimal border states (acting as rewarding distractions), indicating stronger exploratory motives.
Revised: STM agent explores the environment even in the presence of the ubiquitous suboptimal border states (acting as rewarding distractions), indicating stronger exploratory motives.

Original: 6. Conclusion
Revised: 6. Conclusion

Original: We proposed STM, a novel approach to learning a transition dynamics model with the use of sequences of episodic memories, which deﬁne an agent’s more useful, subjective timescale.
Revised: We proposed STM, a novel approach to learning a transition dynamics model with the use of sequences of episodic memories, which deﬁne an agent’s more useful, subjective timescale.

Original: As shown, the emergent properties of STM can be useful to different model-based agents that rely on the transition dynamics model in either training or inference.
Revised: As shown, the emergent properties of STM can be useful to different model-based agents that rely on the transition dynamics model in either training or inference.

Original: We demonstrated its effectiveness within two types of agents and by using several complex 3D foraging tasks within DMLab and AAI.
Revised: We demonstrated its effectiveness within two types of agents and by using several complex 3D foraging tasks within DMLab and AAI.

Original: Nevertheless, a limitation of the current version of STM is its action summarisation heuristic.
Revised: Nevertheless, a limitation of the current version of STM is its action summarisation heuristic.

Original: We pose that integrating general techniques for summarising sequences of actions will allow for more advanced planning and application to environments with more demanding control requirements.
Revised: We pose that integrating general techniques for summarising sequences of actions will allow for more advanced planning and application to environments with more demanding control requirements.

Original: A. Preliminaries A.1. Active Inference
Revised: A. Preliminaries A.1. Active Inference

Original: Active inference is a corollary of the free-energy principle applied to action (Friston et al., 2016; Friston, 2019; Sajid et al., 2019).
Revised: Active inference is a corollary of the free-energy principle applied to action (Friston et al., 2016; Friston, 2019; Sajid et al., 2019).

Original: In this framework, an agent embedded in an environment aims to do two things: (i) minimise surprisal from the observations of the environment under the agent’s internal model of this environment, and (ii) perform actions so as to minimise the expected surprisal in the future.
Revised: In this framework, an agent embedded in an environment aims to do two things: (i) minimise surprisal from the observations of the environment under the agent’s internal model of this environment, and (ii) perform actions so as to minimise the expected surprisal in the future.

Original: More formally, an agent is equipped with a generative model p ( o t , s t ; θ ) , where o t is the agent’s observation at time t , s t is the hidden state of the environment, and θ denotes the parameters of the generative model.
Revised: More formally, an agent is equipped with a generative model p ( o t , s t ; θ ) , where o t is the agent’s observation at time t , s t is the hidden state of the environment, and θ denotes the parameters of the generative model.

Original: The agent’s surprise at time t is deﬁned as the negative log-likelihood, − log p ( o t ; θ ) .
Revised: The agent’s surprise at time t is deﬁned as the negative log-likelihood, − log p ( o t ; θ ) .

Original: We can upper-bound this intractable expression using variational inference by introducing an approximate posterior distribution, q ( s t ) , over s t , such that:
Revised: We can upper-bound this intractable expression using variational inference by introducing an approximate posterior distribution, q ( s t ) , over s t , such that:

Original: [Equation]
Revised: [Equation]

Original: where F is the variational free energy .
Revised: where F is the variational free energy .

Original: The minimisation of this quantity realises objective (i) and is performed by optimising the parameters of the generative model, θ .
Revised: The minimisation of this quantity realises objective (i) and is performed by optimising the parameters of the generative model, θ .

Original: It is also equivalent to the maximisation of model evidence, which intuitively implies that the agent aims to perfect its generative model at explaining the sensory observations from the environment.
Revised: It is also equivalent to the maximisation of model evidence, which intuitively implies that the agent aims to perfect its generative model at explaining the sensory observations from the environment.

Original: 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 :
Revised: 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 :

Original: [Equation]
Revised: [Equation]

Original: where k > t and π = { a t , a t +1 , ..., a k − 1 } is a sequence of actions (policy) between the present time t and the future time k .
Revised: where k > t and π = { a t , a t +1 , ..., a k − 1 } is a sequence of actions (policy) between the present time t and the future time k .

Original: 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.
Revised: 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.

Original: 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:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: where G ( π ) = (cid:80) k>t G ( π, k ) , γ is the temperature parameter, σ ( · ) denotes a softmax function, and t is the present timestep.
Revised: where G ( π ) = (cid:80) k>t G ( π, k ) , γ is the temperature parameter, σ ( · ) denotes a softmax function, and t is the present timestep.

Original: A.2. Dreamer
Revised: A.2. Dreamer

Original: As explained in the main body of the paper, Dreamer operates by learning a policy function q ( a t
Revised: As explained in the main body of the paper, Dreamer operates by learning a policy function q ( a t

Original: | s t ; φ ) using latent roll-outs.
Revised: | s t ; φ ) using latent roll-outs.

Original: While the generative model ( p θ ) is trained to perform accurate reconstructions of the future, the action ( q φ ) and value ( v ψ ) models are trained to maximise rewards (Eq.
Revised: While the generative model ( p θ ) is trained to perform accurate reconstructions of the future, the action ( q φ ) and value ( v ψ ) models are trained to maximise rewards (Eq.

Original: 1) and predict state values (Eq. 2).
Revised: 1) and predict state values (Eq. 2).

Original: As per the original paper by Hafner et al.
Revised: As per the original paper by Hafner et al.

Original: (2020), we outline how the value estimate V λ is calculated.
Revised: (2020), we outline how the value estimate V λ is calculated.

Original: Speciﬁcally, given an imagined trajectory of length H , the value estimate is deﬁned such that:[Equation] = E
Revised: Speciﬁcally, given an imagined trajectory of length H , the value estimate is deﬁned such that:

Original: [Equation]
Revised: [Equation]

Original: (cid:104) h

Original: [Equation]

Original: (cid:88) n =

Original: [Equation]

Original: γ n [Equation] r n + γ h [Equation]
Revised: where h = min ( j + k, t + H ) ,

Original: v

Original: [Equation]
Revised: [Equation]

Original: ( s h ) (cid:105) , where h = min ( j + k, t + H ) , (9)

Original: V λ ( s j ) =

Original: E p θ ,

Original: [Equation]

Original: (cid:104)

Original: [Equation]
Revised: [Equation]

Original: λ )

Original: H

Original: [Equation] 1 (cid:88) n [Equation] (10) where j denotes the timestep from which the roll-out is performed.
Revised: ) (cid:105) , (10) where j denotes the timestep from which the roll-out is performed.

Original: B. Architectural Details and Training B.1. Deep active inference (DAI)
Revised: B. Architectural Details and Training B.1. Deep active inference (DAI)

Original: Baseline system In order to ensure the strongest performance of the baseline agent, we have made slight adjustments to the architecture and training procedure of the original DAI agent by Fountas et al.
Revised: Baseline system In order to ensure the strongest performance of the baseline agent, we have made slight adjustments to the architecture and training procedure of the original DAI agent by Fountas et al.

Original: (2020a), which we report next.
Revised: (2020a), which we report next.

Original: In DAI, each component of the generative and inference models is parametrised by feed-forward neural networks (including fully-connected, convolutional and transpose-convolutional layers), whose architectural details can be found in Table 2.
Revised: In DAI, each component of the generative and inference models is parametrised by feed-forward neural networks (including fully-connected, convolutional and transpose-convolutional layers), whose architectural details can be found in Table 2.

Original: The latent bottleneck of the autoencoder, s , was of size 10.
Revised: The latent bottleneck of the autoencoder, s , was of size 10.

Original: Similarly, we restricted the action space to 3 – forward, left, right.
Revised: Similarly, we restricted the action space to 3 – forward, left, right.

Original: The networks were trained using separate Adam optimisers (Kingma & Ba, 2015) for stability reasons.
Revised: The networks were trained using separate Adam optimisers (Kingma & Ba, 2015) for stability reasons.

Original: The habitual and transition networks are trained with a learning rate of 0 . 0001 ; the autoencoder’s optimiser had a learning rate of 0 . 001 .
Revised: The habitual and transition networks are trained with a learning rate of 0 . 0001 ; the autoencoder’s optimiser had a learning rate of 0 . 001 .

Original: The batch size was set to 50 and the model was trained for 1 × 10 6 iterations under a green observational prior.
Revised: The batch size was set to 50 and the model was trained for 1 × 10 6 iterations under a green observational prior.

Original: All of the networks were implemented using Tensorﬂow v2.2 (Abadi et al., 2015).
Revised: All of the networks were implemented using Tensorﬂow v2.2 (Abadi et al., 2015).

Original: Tests were performed in Animal-AI v2.0.1 (Beyret et al., 2019).
Revised: Tests were performed in Animal-AI v2.0.1 (Beyret et al., 2019).

Original: For the objective-timescale version of DAI, we also employ a top-down attention mechanism, as outlined in Fountas et al.
Revised: For the objective-timescale version of DAI, we also employ a top-down attention mechanism, as outlined in Fountas et al.

Original: (2020a).
Revised: (2020a).

Original: Introducing variable ω that modulates uncertainty about hidden states is aimed at promoting latent state disentanglement and more efﬁcient learning.
Revised: Introducing variable ω that modulates uncertainty about hidden states is aimed at promoting latent state disentanglement and more efﬁcient learning.

Original: Specifically, the latent state distribution is deﬁned as a Gaussian such that s ∼ N ( s ; µ s , σ s /ω ) , where µ s and σ s are the mean and the diagonal covariance, and ω is a decreasing logistic function over the divergence D KL [ q ( a ; φ a ) (cid:107) p ( a )] .
Revised: Specifically, the latent state distribution is deﬁned as a Gaussian such that s ∼ N ( s ; µ s , σ s /ω ) , where µ s and σ s are the mean and the diagonal covariance, and ω is a decreasing logistic function over the divergence D KL [ q ( a ; φ a ) (cid:107) p ( a )] .

Original: During training, the hyperparameters of the top-down attention mechanism were: a = 2 , b = 0 .
Revised: During training, the hyperparameters of the top-down attention mechanism were: a = 2 , b = 0 .

Original: 5 , c = 0 .
Revised: 5 , c = 0 .

Original: 1 , and d = 5 , chosen to match those in Fountas et al.
Revised: 1 , and d = 5 , chosen to match those in Fountas et al.

Original: (2020a).
Revised: (2020a).

Original: Each network was trained with its corresponding loss function, which are the constituent parts of the total variational free energy.
Revised: Each network was trained with its corresponding loss function, which are the constituent parts of the total variational free energy.

Original: In particular, the autoencoder was trained using
Revised: In particular, the autoencoder was trained using

Original: Eqs.
Revised: Eqs.

Original: 11a and 11b, transition model using Eq. 11b, and habitual network using Eq.
Revised: 11a and 11b, transition model using Eq. 11b, and habitual network using Eq.

Original: 11c:
Revised: 11c:

Original: [Equation]
Revised: [Equation]

Original: 11c [Equation] Furthermore, following the training procedure from Fountas et al. (2020a), we stabilise the convergence of the autoencoder by modifying the loss function to:
Revised: Furthermore, following the training procedure from Fountas et al. (2020a), we stabilise the convergence of the autoencoder by modifying the loss function to:

Original: [Equation]
Revised: [Equation]

Original: where γ is a hyperparameter that gradually increases fromto 0.8 during training.
Revised: where γ is a hyperparameter that gradually increases fromto 0.8 during training.

Original: As part of the system’s training procedure, we used a replay buffer with random sampling to mitigate the detrimental effects of on-line learning.
Revised: As part of the system’s training procedure, we used a replay buffer with random sampling to mitigate the detrimental effects of on-line learning.

Original: Full training procedure is shown in Algorithm 2.
Revised: Full training procedure is shown in Algorithm 2.

Original: Adding STM
Revised: Adding STM

Original: The STM module replaces the Markovian transition model of the DAI agent.
Revised: The STM module replaces the Markovian transition model of the DAI agent.

Original: We train the subjective transition model using batch size 15 and a learning rate of 0 . 001 .
Revised: We train the subjective transition model using batch size 15 and a learning rate of 0 . 001 .

Original: Each batch consisted of zero-padded S-sequences with length 50.
Revised: Each batch consisted of zero-padded S-sequences with length 50.

Original: We use a Masking layer to ignore zeropadded parts of the sequences in the computational graph.
Revised: We use a Masking layer to ignore zeropadded parts of the sequences in the computational graph.

Original: The training was stopped at approximately 800 × 10 3 training iterations.
Revised: The training was stopped at approximately 800 × 10 3 training iterations.

Original: The memory formation threshold, (cid:15) , was set to increase via (cid:15) =
Revised: The memory formation threshold, (cid:15) , was set to increase via (cid:15) =

Original: (cid:15) max (1 − exp( − λ · i )) , where threshold ceiling is (cid:15) max = 5 , i is training iteration, λ
Revised: (cid:15) max (1 − exp( − λ · i )) , where threshold ceiling is (cid:15) max = 5 , i is training iteration, λ

Original: = ln 2 /t 1 /with half-life t 1 / 2 = 1 × 10 5 training iterations.
Revised: = ln 2 /t 1 /with half-life t 1 / 2 = 1 × 10 5 training iterations.

Original: B.2. Dreamer
Revised: B.2. Dreamer

Original: Baseline system For the Dreamer agent, we follow the training details outlined in Hafner et al.
Revised: Baseline system For the Dreamer agent, we follow the training details outlined in Hafner et al.

Original: (2020), in order to ensure the strongest performance of the objective-timescale baseline.
Revised: (2020), in order to ensure the strongest performance of the objective-timescale baseline.

Original: The only minor adjustment included setting action repetition to 1, thus forcing the agent to operate over the objective timescale of the environment.
Revised: The only minor adjustment included setting action repetition to 1, thus forcing the agent to operate over the objective timescale of the environment.

Original: The baseline Dreamer was trained for 1.8M steps in DMLab Sparse Rewards and 2.8M steps in DMLab Maze Rewards.
Revised: The baseline Dreamer was trained for 1.8M steps in DMLab Sparse Rewards and 2.8M steps in DMLab Maze Rewards.

Original: Algorithm
Revised: Algorithm

Original: Initialise :
Revised: Initialise :

Original: Empty replay buffer, D ← {}
Revised: Empty replay buffer, D ← {}

Original: Random generative and inference model parameters, { θ, φ } for episode = 1 to ∞ do
Revised: Random generative and inference model parameters, { θ, φ } for episode = 1 to ∞ do

Original: Reset and randomise the environment for t = 1 to 250 do
Revised: Reset and randomise the environment for t = 1 to 250 do

Original: Retrieve observation: ˜ o t
Revised: Retrieve observation: ˜ o t

Original: Sample and take an action ˜ a t ∼ p ( a t ) using the planner
Revised: Sample and take an action ˜ a t ∼ p ( a t ) using the planner

Original: Retrieve next observation: ˜ o t +
Revised: Retrieve next observation: ˜ o t +

Original: Add to replay buffer, D ← D ∪ { (˜ o t , ˜ a t , ˜ o t +1 ) } if D sufﬁcient size then
Revised: Add to replay buffer, D ← D ∪ { (˜ o t , ˜ a t , ˜ o t +1 ) } if D sufﬁcient size then

Original: Sample a batch of transitions: B i ∼ D
Revised: Sample a batch of transitions: B i ∼ D

Original: For every tuple in B i , encode the preceding and the subsequent observations, (˜ o t 1 , ˜ o t 2 ) →
Revised: For every tuple in B i , encode the preceding and the subsequent observations, (˜ o t 1 , ˜ o t 2 ) →

Original: ( q ( s t 1 ; φ s ) , q ( s t 2 ; φ s ))
Revised: ( q ( s t 1 ; φ s ) , q ( s t 2 ; φ s ))

Original: Compute q ( a t 1 | s t 1 ; φ a ) via sampled ˜ s t 1 ∼ q ( s t 1 ; φ s )
Revised: Compute q ( a t 1 | s t 1 ; φ a ) via sampled ˜ s t 1 ∼ q ( s t 1 ; φ s )

Original: Calculate p ( a t 1 ) using the planner
Revised: Calculate p ( a t 1 ) using the planner

Original: Train habitual network (cid:46) Eq.
Revised: Train habitual network (cid:46) Eq.

Original: 11c
Revised: 11c

Original: Compute ω t 2 with the use of Eq.
Revised: Compute ω t 2 with the use of Eq.

Original: 11c
Revised: 11c

Original: Compute next state sufﬁcient statistics ( µ t 2 , σ t 2 ) using p ( s t 2 | s t 1 , a t 1 ; θ s )
Revised: Compute next state sufﬁcient statistics ( µ t 2 , σ t 2 ) using p ( s t 2 | s t 1 , a t 1 ; θ s )

Original:  (cid:46) Eq. 11b
Revised: Train transition model (cid:46) Eq. 11b

Original: Train variational autoencoder (cid:46) Eqs.
Revised: Train variational autoencoder (cid:46) Eqs.

Original: 11a, 11b
Revised: 11a, 11b

Original: Adding STM Because the Dreamer already involved a recurrent dynamics model, no changes to the architecture of the Dreamer+STM agent needed to be done.
Revised: Adding STM Because the Dreamer already involved a recurrent dynamics model, no changes to the architecture of the Dreamer+STM agent needed to be done.

Original: As such, the training procedure with all concomitant hyperparameters were kept the same.
Revised: As such, the training procedure with all concomitant hyperparameters were kept the same.

Original: The threshold was set to a constant value, (cid:15)
Revised: The threshold was set to a constant value, (cid:15)

Original: =
Revised: =

Original: 7 for Sparse Rewards and (cid:15) = 8 for Maze
Revised: 7 for Sparse Rewards and (cid:15) = 8 for Maze

Original: Rewards.
Revised: Rewards.

Original: Dreamer+STM was trained for 1.8M steps in
Revised: Dreamer+STM was trained for 1.8M steps in

Original: DMLab Sparse Rewards and 3.5M steps in DMLab Maze
Revised: DMLab Sparse Rewards and 3.5M steps in DMLab Maze

Original: Rewards.
Revised: Rewards.

Original: B.3. Action Heuristic
Revised: B.3. Action Heuristic

Original: To train the STM transition model in the AAI and DMLab environments, we implement a simple heuristic that is used to summarise a sequence of actions taken by the agent from one memory to reach the next one.
Revised: To train the STM transition model in the AAI and DMLab environments, we implement a simple heuristic that is used to summarise a sequence of actions taken by the agent from one memory to reach the next one.

Original: A sequence of actions, A = { a τ 1 , a τ 1 +1 , ...a τ 1 +( N − 1) } , takes the agent from a recorded memory s τ 1 to memory s τ 2 , where the time between these states τ 2 − τ 1 = N , and a ∈ { a forward , a right , a left } .
Revised: A sequence of actions, A = { a τ 1 , a τ 1 +1 , ...a τ 1 +( N − 1) } , takes the agent from a recorded memory s τ 1 to memory s τ 2 , where the time between these states τ 2 − τ 1 = N , and a ∈ { a forward , a right , a left } .

Original: We employ polar coordinates relative to the agent’s initial position in Cartesian coordinates at time τ 1 and perform iterative updates of its position after every action until the time-step of the next episodic memory,[Table] [Table]
Revised: We employ polar coordinates relative to the agent’s initial position in Cartesian coordinates at time τ 1 and perform iterative updates of its position after every action until the time-step of the next episodic memory,

Original: [Table]
Revised: [Table]

Original: Sigmoid

Original: STM Transition

Original: Linear

Original: ReLU LSTMLSTM - LSTMLSTM - Linearτ 2 , is reached.
Revised: τ 2 , is reached.

Original: Given the agent’s orientation in the environment, θ , the next position of the agent given a forward action is calculated using, p [Equation] sin [Equation]
Revised: Given the agent’s orientation in the environment, θ , the next position of the agent given a forward action is calculated using, p t +1 =

Original:  cos [Equation] , wherep [Equation]
Revised: p t + (cid:20) sin θ cos θ (cid:21) , where

Original: =

Original: [Equation]
Revised: [Equation]

Original: Finally, we retrieve angle φ , which describes the direction in which the agent has travelled with respect to its initial position and orientation.
Revised: Finally, we retrieve angle φ , which describes the direction in which the agent has travelled with respect to its initial position and orientation.

Original: This angle is used to decide on the action that summarises the trajectory (see Figure 8) using 
Revised: This angle is used to decide on the action that summarises the trajectory (see Figure 8) using a =   a forward |

Original: [Equation]
Revised: [Equation]

Revised: 5 ◦ a right 22 .

Revised: 5 ◦ < φ < 180 ◦ a left − 22 .

Revised: 5 ◦ > φ ≥ − 180 ◦ .

Original: Although this heuristic provided satisfactory results, trajectory encoding is one of the most limiting parts of the
Revised: Although this heuristic provided satisfactory results, trajectory encoding is one of the most limiting parts of the

Original: STM.
Revised: STM.

Original: Nevertheless, we believe such simplistic encoding was sufﬁcient to demonstrate the primary characteristics of the STM – relating to time and prediction saliency.
Revised: Nevertheless, we believe such simplistic encoding was sufﬁcient to demonstrate the primary characteristics of the STM – relating to time and prediction saliency.

Original: As such, we strongly encourage further research into more sophisticated and generalisable techniques of summarising action sequences.
Revised: As such, we strongly encourage further research into more sophisticated and generalisable techniques of summarising action sequences.

Original: C. Testing C.1. Animal-AI Environment
Revised: C. Testing C.1. Animal-AI Environment

Original: We expect some readers to be unfamiliar with the AnimalAI environment (Beyret et al., 2019), and thus provide a few snapshots from the environment in Figure 9.
Revised: We expect some readers to be unfamiliar with the AnimalAI environment (Beyret et al., 2019), and thus provide a few snapshots from the environment in Figure 9.

Original: C.2. Monte Carlo Tree Search of DAI
Revised: C.2. Monte Carlo Tree Search of DAI

Original: Following Fountas et al.
Revised: Following Fountas et al.

Original: (2020a), we deﬁne the MCTS upper conﬁdence bound as,
Revised: (2020a), we deﬁne the MCTS upper conﬁdence bound as,

Original: [Equation]
Revised: [Equation]

Original: where ˜ G ( s t , a t ) is agent’s estimation of the expected free energy given action a t from state s t , N ( a t , s t ) is the number of times a t was taken from s t during tree search, and c explore is a tree exploration hyperparameter.
Revised: where ˜ G ( s t , a t ) is agent’s estimation of the expected free energy given action a t from state s t , N ( a t , s t ) is the number of times a t was taken from s t during tree search, and c explore is a tree exploration hyperparameter.

Original: For testing, the MCTS parameters of the baseline and STM were kept the same to ensure a fair comparison, with the exception of the exploration hyperparameter, c explore , which was optimised separately for both to yield the best performance of the agents: c explore = 1 for STM and c explore =for the baseline.
Revised: For testing, the MCTS parameters of the baseline and STM were kept the same to ensure a fair comparison, with the exception of the exploration hyperparameter, c explore , which was optimised separately for both to yield the best performance of the agents: c explore = 1 for STM and c explore =for the baseline.

Original: Otherwise, both agents used 15 MCTS simulation loops with depth 3.
Revised: Otherwise, both agents used 15 MCTS simulation loops with depth 3.

Original: D. Empirical Trends
Revised: D. Empirical Trends

Original: As discussed in Section 4.1, prediction error (or free energy) of the transition model was used as the criterion for memory formation.
Revised: As discussed in Section 4.1, prediction error (or free energy) of the transition model was used as the criterion for memory formation.

Original: Although the idea is largely inspired by the neuroscience literature on episodic memory and human time perception, it is useful to examine the practical implications of this criterion in the context of model-based agents.
Revised: Although the idea is largely inspired by the neuroscience literature on episodic memory and human time perception, it is useful to examine the practical implications of this criterion in the context of model-based agents.

Original: In particular, we consider the necessity for context-aware adaptation of the temporal extent of future predictions during planning.
Revised: In particular, we consider the necessity for context-aware adaptation of the temporal extent of future predictions during planning.

Original: We pose that using shorter extents in perceptually-complex scenes is conducive to the agent’s ability to navigate difﬁcult or important parts of the environment.
Revised: We pose that using shorter extents in perceptually-complex scenes is conducive to the agent’s ability to navigate difﬁcult or important parts of the environment.

Original: In contrast, longer extents should be used in scenes with little perceptual complexity, as the agent plans to get to the objects of its interest.
Revised: In contrast, longer extents should be used in scenes with little perceptual complexity, as the agent plans to get to the objects of its interest.

Original: This is further justiﬁed by the STM agent’s signiﬁcantly improved performance at avoiding negative rewards.
Revised: This is further justiﬁed by the STM agent’s signiﬁcantly improved performance at avoiding negative rewards.

Original: The chosen criterion allows the STM transition model to learn just that – due to the fact that the frequency of memory formation increases in more complex or rare scenes, as the prediction error gets larger.
Revised: The chosen criterion allows the STM transition model to learn just that – due to the fact that the frequency of memory formation increases in more complex or rare scenes, as the prediction error gets larger.

Original: Empirically, this is demonstrated in Figure 10, in which observations from a replay buffer were sorted based on the DAI agent’s free energy of the transition model.
Revised: Empirically, this is demonstrated in Figure 10, in which observations from a replay buffer were sorted based on the DAI agent’s free energy of the transition model.

Original: It can be seen that observations with higher free energy contain more objects and, generally, constitute more perceptually-complex scenes.
Revised: It can be seen that observations with higher free energy contain more objects and, generally, constitute more perceptually-complex scenes.

Original: At the same time, higher values of transition model free energy can be driven by epistemically-surprising events (e.g. sphere entering the agent’s view upon turning).
Revised: At the same time, higher values of transition model free energy can be driven by epistemically-surprising events (e.g. sphere entering the agent’s view upon turning).

Original: Indeed, these types of observations can also be observed in the higher levels of Figure 10.
Revised: Indeed, these types of observations can also be observed in the higher levels of Figure 10.

Original: As demonstrated in the paper, such observations appearing in the recorded S-sequences drive the STM to producing more salient predictions of the future, which in turn encourages exploratory behaviour.
Revised: As demonstrated in the paper, such observations appearing in the recorded S-sequences drive the STM to producing more salient predictions of the future, which in turn encourages exploratory behaviour.

Original: E. Additional Results
Revised: E. Additional Results

Original: Collected S-sequences Figure 11 shows additional examples of collected S-sequences in DMLab.
Revised: Collected S-sequences Figure 11 shows additional examples of collected S-sequences in DMLab.

Original: In particular, we draw reader’s attention to how the model is able to record memories at different temporal rates, in which different segments are approximately denoted with different colours.
Revised: In particular, we draw reader’s attention to how the model is able to record memories at different temporal rates, in which different segments are approximately denoted with different colours.

Original: In visually simple settings, the model tends to skip more frames than in more complex scenes, thus resulting in the STM’s ability to vary the extent of its predictions.
Revised: In visually simple settings, the model tends to skip more frames than in more complex scenes, thus resulting in the STM’s ability to vary the extent of its predictions.

Original: To better understand STM memory collection, access the following anonymised link, which contains a number of videos of episodes as perceived by the objective- and subjective-timescale agents.
Revised: To better understand STM memory collection, access the following anonymised link, which contains a number of videos of episodes as perceived by the objective- and subjective-timescale agents.

Original: https: //www.notion.so/Dreamer-STM-Episodes47692c7680b541d69e13983d04256aec .
Revised: https: //www.notion.so/Dreamer-STM-Episodes47692c7680b541d69e13983d04256aec .

Original: STM roll-outs Figure 13 shows random STM roll-outs generated by the system.
Revised: STM roll-outs Figure 13 shows random STM roll-outs generated by the system.

Original: These diverse roll-outs demonstrate that STM is able to: i) make correct actionconditioned predictions, ii) speed up its prediction timescale when objects are far away, iii) slow down the prediction timescale when objects are nearby.
Revised: These diverse roll-outs demonstrate that STM is able to: i) make correct actionconditioned predictions, ii) speed up its prediction timescale when objects are far away, iii) slow down the prediction timescale when objects are nearby.

Original: Additionally, for initial observations without any objects in the frame, the agent is capable of imagining their presence in the future.
Revised: Additionally, for initial observations without any objects in the frame, the agent is capable of imagining their presence in the future.

Original: Exploration
Revised: Exploration

Original: heatmaps The provided exploration heatmaps in the main body of the paper were concerned with showcasing the STM’s ability to exert more pronounced exploratory behaviour in the presence of rewarding distractions.
Revised: heatmaps The provided exploration heatmaps in the main body of the paper were concerned with showcasing the STM’s ability to exert more pronounced exploratory behaviour in the presence of rewarding distractions.

Original: Here, we additionally demonstrate the exploratory behaviour of Dreamer with and without STM in an empty DMLab environment (in the absence of any rewards) – see Figure 14.
Revised: Here, we additionally demonstrate the exploratory behaviour of Dreamer with and without STM in an empty DMLab environment (in the absence of any rewards) – see Figure 14.

Original: Speciﬁcally, Figure 14a shows the exploration heatmap comparison, while Figure 14b shows 20 random samples of the different paths that the agents took in an empty arena.
Revised: Speciﬁcally, Figure 14a shows the exploration heatmap comparison, while Figure 14b shows 20 random samples of the different paths that the agents took in an empty arena.

Original: We emphasise the considerable differences in the agents’ behaviours –
Revised: We emphasise the considerable differences in the agents’ behaviours –

Original: STM agent employs a consistent strategy with extensive and persistent exploration of the environment, while the objective-timescale agent resorts to more chaotic behaviour.
Revised: STM agent employs a consistent strategy with extensive and persistent exploration of the environment, while the objective-timescale agent resorts to more chaotic behaviour.

Original: 47692c7680b541d69e13983d04256aec

Original: Firstly, the agent can successfully model the action-conditioned dynamics of the environment, justifying the appropriateness of the used action heuristic.
Revised: Firstly, the agent can successfully model the action-conditioned dynamics of the environment, justifying the appropriateness of the used action heuristic.

Original: Secondly, in line with the analysis from the main body of the paper, STM is able to vary the extent of its predictions depending on the perceptual context.
Revised: Secondly, in line with the analysis from the main body of the paper, STM is able to vary the extent of its predictions depending on the perceptual context.

Original: For instance, the top roll-out illustrates a rapid approach to a red sphere, while slowing down close to it.
Revised: For instance, the top roll-out illustrates a rapid approach to a red sphere, while slowing down close to it.

Original: This is akin to some of the presented S-sequences in Figure 11, which illustrate a very similar behaviour.
Revised: This is akin to some of the presented S-sequences in Figure 11, which illustrate a very similar behaviour.

Original: Thirdly, in roll-outs with uninformative initial states (such as the ones with the agent facing a wall), STM agent tends to imagine observing an object in the future and with generally increased speed of future predictions.
Revised: Thirdly, in roll-outs with uninformative initial states (such as the ones with the agent facing a wall), STM agent tends to imagine observing an object in the future and with generally increased speed of future predictions.

Original: [Figure]
Revised: [Figure]