Machines that can replicate human intelligence with type 2 daniel2017thinking reasoning capabilities should be able to reason at multiple levels of spatio-temporal abstractions and scales using internal world models friston2008hierarchical,lecun2022path. Devising formalisms to develop such internal world models, which accurately reflect the causal hierarchies inherent in the dynamics of the real world, is a critical research challenge in the domains of artificial intelligence and machine learning. This thesis identifies several limitations with the prevalent use of state space models (SSMs) as internal world models and propose two new probabilistic formalisms namely Hidden-Parameter SSMs and Multi-Time Scale SSMs to address these drawbacks. The structure of graphical models in both formalisms facilitates scalable exact probabilistic inference using belief propagation, as well as end-to-end learning via backpropagation through time. This approach permits the development of scalable, adaptive hierarchical world models capable of representing nonstationary dynamics across multiple temporal abstractions and scales. Moreover, these probabilistic formalisms integrate the concept of uncertainty in world states, thus improving the system's capacity to emulate the stochastic nature of the real world and quantify the confidence in its predictions. The thesis also discuss how these formalisms are in line with related neuroscience literature on Bayesian brain hypothesis and predicitive processing. Our experiments on various real and simulated robots demonstrate that our formalisms can match and in many cases exceed the performance of contemporary transformer variants in making long-range future predictions. We conclude the thesis by reflecting on the limitations of our current models and suggesting directions for future research.

Maschinen, die menschliche Intelligenz mit Fähigkeiten zur Schlussfolgerung des Typs 2 daniel2017thinking replizieren können, sollten in der Lage sein, auf mehreren Ebenen räumlich-zeitlicher Abstraktionen und Maßstäbe zu schlussfolgern, indem sie interne Weltmodelle verwenden friston2008hierarchical,lecun2022path. Die Entwicklung von Formalismen zur Entwicklung solcher internen Weltmodelle, die die in der Dynamik der realen Welt inhärenten kausalen Hierarchien genau widerspiegeln, stellt eine entscheidende Forschungsherausforderung in den Bereichen Künstliche Intelligenz und maschinelles Lernen dar. Diese Dissertation identifiziert mehrere Einschränkungen bei der verbreiteten Verwendung von Zustandsraummodellen (SSMs) als interne Weltmodelle und schlägt zwei neue probabilistische Formalismen vor, nämlich Versteckte-Parameter-SSMs und Mehr-Zeitskalen-SSMs, um diese Nachteile anzugehen. Die Struktur der grafischen Modelle in beiden Formalismen ermöglicht skalierbare exakte probabilistische Inferenz mittels Belief Propagation sowie End-to-End-Lernen durch Backpropagation durch die Zeit. Dieser Ansatz ermöglicht die Entwicklung skalierbarer, adaptiver hierarchischer Weltmodelle, die nichtstationäre Dynamiken über mehrere zeitliche Abstraktionen und Maßstäbe darstellen können. Darüber hinaus integrieren diese probabilistischen Formalismen das Konzept der Unsicherheit in Weltzuständen und verbessern somit die Fähigkeit des Systems, die stochastische Natur der realen Welt nachzuahmen und das Vertrauen in seine Vorhersagen zu quantifizieren. Die Arbeit diskutiert auch, wie diese Formalismen mit der verwandten neurowissenschaftlichen Literatur zur Hypothese des Bayesianischen Gehirns und zur prädiktiven Verarbeitung übereinstimmen. Unsere Experimente mit verschiedenen realen und simulierten Robotern zeigen, dass unsere Formalismen in vielen Fällen die Leistung zeitgenössischer Transformer-Varianten bei der Vorhersage langfristiger Zukünfte übertreffen können. Wir schließen die Dissertation mit einer Reflexion über die Grenzen unserer aktuellen Modelle und Vorschlägen für zukünftige Forschungsrichtungen ab.

CHAPTER: ACKNOWLEDGEMENT

This thesis is the result of the support of several people, to whom I am extremely grateful. I would first like to express my thanks and gratitude to my supervisor, Prof. Gerhard Neumann, for providing me the opportunity to work with him and his consistent mentoring over the last 5 years. His mentorship has helped shape my approach to research, introducing me to a probabilistic way of thinking and emphasizing the importance of rigor. He is one of the most approachable people I have met in my life and carries himself with a deep sense of humility and patience, qualities that continue to inspire me both professionally and personally.
I am also thankful for the guidance and support from several other mentors across various projects during my PhD. I extend my gratitude to Professor Marc Hanheide, Dr. Dieter Buchler, Harit Pandya, Philipp Becker, Aravind Srinivas, Ozan Demir and Niels van Duijkeren, who each played a crucial role at different stages of my doctoral journey.
My heartfelt thanks go to our administrative secretary, Christine Brand, for her kind support over the PhD, which made navigating daily challenges as a non-German speaking foreigner much simpler.
One of the highlights of my PhD journey was being able to be part of two separate research groups, at KIT in Germany and the University of Lincoln in the UK. This provided me an incredible opportunity to immerse myself in two distinct cultures. I am indebted to the several brilliant, kind, and inclusive colleagues at both these research groups with whom I had the privilege of sharing many hikes, retreats, ski trips, dinners, coffee breaks, and gossip sessions. The technical debates I had with them have significantly contributed to molding my thought process as a researcher. I would like to extend special thanks to Harit Pandya, Riccardo, Philipp Becker, Onur, Bruce (Ge Li), and Max, whose friendship and support made the initial phases of my PhD journey and the transition between two countries significantly easier.
Special thanks to my friend Sharvin Raphael, who assisted with the Photoshop-based graphics used in this thesis, enhancing the visual presentation of my research.
It was also a privilege to work with and mentor several brilliant students (Andreas Boltres, Martin Herault, Moritz Reuss, Rohit Sonker, Ruben Jacob, Stefan Geyer, Shuheng Zhang, Emiliyan Gospodinov, Thorben Comes, among others). The discussions with them were invaluable, significantly influencing my thought process. Additionally, working with them helped me discover my passion for mentoring and teaching, for which I am immensely grateful.
I extend my heartfelt gratitude to my parents, brother, friends, and family, both in India and overseas, whose constant support and encouragement were crucial to my pursuit and perseverance in this ambitious endeavor of moving to Europe for my PhD. Finally, special thanks to Kevin for making my experience in Germany both rich and colourful.

CHAPTER: INTRODUCTION

The human brain maintains an intricate internal mental model of the world, a sophisticated representation that allows us to navigate complex environments, predict outcomes, and learn from interactions with minimal explicit instruction. This concept of "World Models" sutton1995td,ha2018world,lecun2022path in machine learning research draws direct inspiration from our cognitive processes friston2005theory,friston2008hierarchical, aiming to endow AI systems with a similar ability to abstract, understand, and anticipate the dynamics of their surroundings, a process that remains challenging for current AI systems. Research on foundational world models is pivotal for advancing beyond the current limitations of machine learning, where systems often require vast datasets and extensive training to perform tasks that humans manage with ease.

Recent advances in transformer vaswani2017attention based large language models (LLMs) like BERT devlin2018bert, GPT-3,4 brown2020language,kocon2023chatgpt etc have impressed the AI community due to their remarkable ability to generate coherent text, translate languages, and even produce code. However, the perception that LLMs might be the definitive pathway to AGI overlooks critical nuances. At their core, LLMs excel in pattern recognition and statistical inference from vast datasets but lack an understanding of causality, physical principles, and reasoning—essential components of human-like intelligence. Recent research zevcevic2023causal on the causal nature of LLMs argues that they rely heavily on correlational statistics and fail to capture the causal dynamics essential for understanding and interacting with the real world. By focusing on learning the causal structure of the world and enabling machines to simulate and predict future states, researchers can create more autonomous, efficient, and adaptable AI systems. Such a research path represents a measured step towards developing AI systems with a deeper, more intuitive grasp of the world, aligning machine learning processes more closely with natural intelligence. Advancements in learning a foundational world model hold the promise of profound implications across a myriad of fields, including autonomous vehicles, robotics, healthcare, environmental modelling, and beyond, marking a crucial step towards achieving AI that can truly understand and navigate the world as humans do pezzulo2018hierarchical, lecun2022path.

When discussing foundational world model (FWM), I am referring to a model (or a set of models) that embodies the following characteristics expanding on the definition from  <cit.>.

* (Representation) Conceptually understand the components, structures, and interaction dynamics within a given system at different levels of abstractions lee2003hierarchical,friston2008hierarchical,pezzulo2018hierarchical,lecun2022path;

* (Veridicality) Quantitatively model the underlying laws of such a system, that enables accurate predictions of counterfactual consequences of interventions/actions;

* (Probabilistic) Model uncertainties (aleatoric  and epistemic) via probabilistic representations. This ensures that the FWM can handle the stochastic nature of the real world, quantify confidence in its predictions and imagine diverse plausible futures.

* (Foundational) Being able to generalize and scale (i), (ii) and (iii) across diverse systems or domains encountered in the world.

The definition of FWM in <cit.> is expanded to explicitly include the "probabilistic" property. The reasoning for this is further elaborated in Section <ref>. Furthermore, Section <ref> provides a reasoning for this from a computational neuroscience point of view.

§ THESIS PROBLEM STATEMENT

In this dissertation, I explore the difficulties associated with existing world model formalisms to fit into the notion of FWM and introduce formalism(s) designed to tackle three challenges in developing foundational models of the world, as outlined below.

* How can we model the world with a principled probabilistic formalism that also quantifies uncertainties in predictions in a scalable manner?

The real world involves a significant amount of uncertainty. It is intrinsically stochastic (aleatoric uncertainty), and we are often
uncertain about the true state of the system because our observations about it are partial (epistemic uncertainty), or the world model’s prediction accuracy is imperfect due to limited training data, representational power, or computational constraints. This necessitates probabilistic world models that can imagine multiple plausible future scenarios and the associated uncertainty. Furthermore, the inference algorithms on these probabilistic models/representations should scale well for large datasets and long sequences.

* How can we make sure the model adapts to changing dynamics / non-stationary situations?
In various real-world control scenarios, agents are presented with tasks that may exhibit different dynamics due to changing environmental factors or task requirements. For instance, a robot engaged in playing table tennis might use rackets of differing weights or sizes, or an agent tasked with bottle manipulation might deal with varying fluid levels. Traditional models, operating on a singular time scale or with flat world representations, struggle to accommodate these dynamic shifts. Therefore, it's essential to develop formalisms capable of inferring task abstractions for seamless adaptability in multi-task environments.

* Can we model the world at multiple temporal abstractions and time scales?

One important dimension of world models is the level of temporal granularity/time scale at which they operate. Existing literature on world models operates at a single level of temporal abstraction, typically at a fine-grained level such as milliseconds. One drawback of single-time scale world models is that they may not capture longer-term trends and patterns in the data. For efficient long-horizon prediction, reasoning and planning, the model needs to predict at multiple levels of temporal abstractions.

The challenges outlined here are not exhaustive of all the difficulties encountered in developing foundational world models (FWMs). Instead, they specifically target three central aspects that are vital for the successful creation of FWMs.

§ THESIS CONTRIBUTION

This thesis introduces a systematic approach to address the said challenges in a principled probabilistic framework through the development of three formalisms, each building upon the other. Each formalism exhibits increasing complexity and expressiveness, with the incorporation of one utility at a time, culminating in a final model that combines the favourable properties of all three. The sequential presentation also ensures a comprehensive understanding of the progression from the foundational state space models (SSMs) to the sophisticated Multi-time scale SSMs, showcasing a systematic evolution in complexity and capability.

Evolution of thesis represented in terms of the evolution of the probabilistic graphical models for each world model formalism. All formalisms infer the latent states via exact inference resulting in closed-form update rules.

The main contributions that culminate from this thesis are (i) identifying the drawbacks of the widely used SSM formalism and (ii) introducing two new formalisms / declarative representations koller2009probabilistic called HiP-SSM and Multi-Time Scale SSM to address these drawbacks. The thesis also proposes algorithms to perform scalable exact inference via message passing for forward predictions in these models. Finally, it is demonstrated that end-to-end learning is possible in all of the proposed models using a single loss function and backpropagation through time.

§ THESIS OUTLINE

The remainder of this thesis is structured as follows. A visual overview of the thesis is given in Figure <ref>.

In Chapter 2, we give a list of preliminaries that will assist readers in navigating the rest of the document. This chapter begins with a concise introduction to Probabilistic Graphical Models (PGMs), a core framework utilized throughout this thesis for formalizing representations of the world dynamics and performing learning and inference in these. Additionally, we introduce fundamental concepts, including deep state space models (SSMs) and Bayesian aggregation schemes.

A visual overview of the thesis.

Chapter 3 reviews the existing literature on world models, drawing from two distinct areas: computational neuroscience and machine learning. The segment on computational neuroscience aims to highlight connections between the representations explored in this thesis and the conceptualizations of world models in the human brain set forth by the neuroscience community. Additionally, we succinctly review various methodologies adopted by the machine learning community for building world models, providing a backdrop for the thesis's further discussions.

In Chapter 4, we use the well-established formalism of state space models (SSMs) for learning single time-scale World Models. We critically analyze and extend a recent deep Kalman model (deep Gaussian SSM) known as recurrent Kalman networks (RKNs) becker2019recurrent for the specific objective of modeling world dynamics. A significant limitation identified in the original models was their inability to systematically include control or action input in the latent dynamics. This chapter suggests a methodologically sound approach to embed action signals into state space models that respect the causal relation with the inferred latent states. Conditioning actions in a manner that adheres to causal relationships is essential, particularly for decision-making and control tasks. This approach enables precise execution of interventions ("doing") and the simulation of counterfactual scenarios ("imagining") using actions or control signals. The hypothesis is validated by modelling the dynamics of a variety of real robots with non-markovian but stationary dynamics.

In Chapter 5, we focus on the aspect of "adaptability" of world models within dynamic, non-stationary environments. Traditional State Space Models (SSMs) typically operate under the assumption of static dynamics, a presumption that falls short in the complexity of real-world applications. Recognizing the challenge of developing distinct SSMs for all possible variation in dynamics, we introduce the formalism of Hidden Parameter State Space Models (HiP-SSMs). HiP-SSMs leverage hierarchical latent task abstractions to infer the causal factors behind environmental non-stationarity, facilitating seamless adaptation to changing conditions. Moreover, we detail a simple and computationally efficient method for learning and inference in this Gaussian graphical model that avoids approximations like variational inference.

In Chapter 6, the thesis explores the notion of "temporal depth" in the world dynamics via causal hierarchies. Learning world models with hierarchical temporal abstractions and multiple time scales is a largely unaddressed problem in existing literature for world models (deterministic or stochastic). We come up with a principled probabilistic formalism for learning such multi-time scale world models as a hierarchical sequential latent variable model called multi-time scale state space models (MTS3). MTS3 maintains all the advantages of the previous formalisms (SSM and HiP-SSM)  for learning world models but can additionally make hierarchical long-horizon predictions. We also derive computationally efficient inference schemes on multiple time scales for highly accurate long-horizon predictions and uncertainty estimates spanning several seconds into the future. We show these lightweight hierarchical linear models can compete with state of the art transformers on long-horizon predictions on continuous dynamical systems.

Chapter 7 concludes the thesis by summarizing the findings, outlining the limitations, and directions for future research.

CHAPTER: PRELIMINARIES

§ PROBABILISTIC GRAPHICAL MODELS

Probabilistic Graphical Models (PGMs) jordan2004graphical,koller2007graphical use a graph-based representation as the basis for compactly encoding a complex distribution over a high-dimensional space. In various applied fields including bioinformatics, speech processing, image processing and control theory, statistical models have long been formulated in terms of graphs, and algorithms for computing basic statistical quantities such as likelihoods and score functions have often been expressed in terms of recursions operating on these graphs; examples include phylogenies, pedigrees, hidden Markov models, Markov random fields, and Kalman filters. These ideas can be understood, unified, and generalized within the formalism of graphical models. Indeed, graphical models provide a natural tool for formulating variations on these classical architectures, as well as for exploring entirely new families of statistical models.

The intuitive formalisms via PGMs enable effective communication between scientists across the mathematical divide by fostering substantive debate in the context of a scientific problem and ultimately facilitate the joint development of statistical and computational tools for quantitative data analysis. Accordingly, in fields that involve the study of large numbers of interacting variables, graphical models are increasingly in evidence.

§.§ Structure Of Graph and Independencies

There is a dual perspective that can be used to interpret the structure of the graph in PGMs.

* The graph as a representation of a set of
independencies: From one perspective, the graph is a compact representation of a set of independencies that hold in the distribution; these properties take the form $X$ is independent of $\boldsymbol{Y}$ given $Z$, denoted $\boldsymbol{X} \perp \boldsymbol{Y} \mid \boldsymbol{Z})$, for some subsets of variables $\boldsymbol{X}, \boldsymbol{Y}, \boldsymbol{Z}$.

* The graph as a skeleton for factorizing a distribution: The other perspective is that the graph defines a skeleton for compactly representing a high-dimensional distribution: Rather than encode the probability of every possible assignment to all of the variables in our domain, we can "break up" the distribution into smaller factors, each over a much smaller space of possibilities.

The two perspectives are, in a
deep sense, equivalent. The independence properties of the distribution are precisely what allow it to be represented compactly in a factorized form. Conversely, a particular
factorization of the distribution guarantees that certain independencies hold.

(left) A Bayesian Network, where the directed edges indicate causal relationships. (right) A Markov Network, where the undirected edges give a notion of correlation/affinity.

§.§ Representation, Inference and Learning in PGMs

§.§.§ Representation
As detailed in section <ref>, the representation of distributions in graphical language exploits structure that appears in many distributions that we want to encode in practice. The graph structure allows us to factorize a high-dimensional joint distribution into a product of lower-dimensional CPDs or factors exploiting conditional independencies. This framework has many advantages as listed below.

* It often allows the distribution to be written down tractably, even in cases where the explicit representation of the joint distribution is astronomically large.

* A human expert can understand and evaluate its semantics and properties. This property is important for constructing models that provide an accurate reflection of our understanding of a domain.

PGMs can be broadly classified into 2 families based on the types of their graphical structure, Bayesian and Markov networks. Bayesian Networks uses a directed graph that captures causal relationships between variables. Markov Networks uses an undirected graph and captures correlations between variables. Both representations provide the duality of independencies and factorization, but they differ in the set of independencies they can encode and in the factorization of the distribution that they induce.

In this thesis, we come up with formalisms for internal mental models of the world/environment dynamics using these graph-based representations, specifically Bayesian Networks.

§.§.§ Inference
The same structure often also allows the distribution to be used effectively for inference — answering queries using the distribution as our model of the world. In particular, there exist algorithms for computing the posterior probability of some variables given evidence of others. These inference algorithms work directly on the graph structure and are generally orders of magnitude faster than manipulating the joint distribution explicitly.

A summary of standard inference routines on PGMs.

This thesis perform exact inference via closed form messages (belief propagation) to answer several such queries on the proposed graphical representation of the world model.

§.§.§ Learning
PGM framework facilitates the effective construction of approximate models by learning from past experience (in the form of observed data). In this approach, a human expert provides some rough guidelines on how to model a given domain. For example, the human usually specifies the attributes that the model should contain, often some of the main dependencies that it should encode, and perhaps other aspects. However, the details are usually filled in automatically by fitting the model to the data. The models produced by this process are usually much better reflections of the domain than models that are purely hand-constructed. In addition, they can sometimes reveal surprising connections between variables and provide novel insights into a domain. Learning in PGMs can be broadly classified into (i) structure learning and (ii) parameter learning. Structure learning focuses on discovering the underlying causal or dependence structure between the variables in the model. It essentially answers the question "What variables influence each other?". Parameter learning, on the other hand, focuses on estimating the quantitative relationships between variables based on the discovered structure. It essentially answers the question "How strong are these relationships?".

This thesis proposes end-to-end parameter learning strategies for the Bayesian Network representations from data via backpropagation.

§.§ Predictive World Models As Inference In PGMs

In this thesis, we use PGM as a tool to model the causal dynamics of the world in a probabilistically principled manner. Indeed, if a particular learning problem can be set up as a probabilistic graphical model, this can often serve as the first and most important step in solving it. The PGM framework offers us the flexibility that once we write down the model and pose the question, the objectives for learning and inference emerge automatically.

To set up models that reflect the causal structure of our world/environment dynamics, we choose directed graphical models (Bayesian Networks) as declarative graphical representations. We can then effectively use the representation to answer a wide range of questions that are of interest using well-studied and efficient inference algorithms. These queries may include (i) "state queries" which seek answers about the past, present and future world states, and (ii) "causal queries" like intervention and counterfactual queries. In this thesis, we perform both "state queries" and "causal queries". The "causal queries" primarily take the form of action-conditional future predictions. Finally, for those parameters that are unknown to the human expert, principled learning objectives can be used to learn parameters from observed data. A summary of the schemes utilized in the thesis is given in Table <ref>.

We start with the most widely used and studied PGM for modeling the world, commonly referred to as State Space Models (SSMs). This formalism will form the basis for all other formalisms developed in this thesis.

Types     Used In Thesis
Representation     Bayesian Network     $\checkmark$
Markov Network
Inference     Exact Inference     $\checkmark$

Approximate Inference
Learning     Parameter Learning    $\checkmark$

Structural Learning

Backpropagation (Gradient Descent)    $\checkmark$

Expectation Maximization

A summary of various representation, learning and inference approaches used in the Thesis.

§ STATE SPACE MODELS

PGM for a State Space Model (SSM).

State space models (SSMs) are Bayesian probabilistic graphical models koller2009probabilistic, jordan2004graphical that are popular for learning patterns and predicting behaviour in sequential data and dynamical systems.

Formally, we define a state space model as a tuple $( \mathcal{Z}, \mathcal{A}, \mathcal{O}, f, h, \Delta t)$, where $\mathcal{Z}$ is the state space, $\mathcal{A}$ the action space and $\mathcal{O}$ the observation space of the SSM. The parameter $\Delta t$ denotes the discretization time step and $f$ and $h$ the dynamics and observation models, respectively.

We will consider the Gaussian state space model that is represented using the following equations

\begin{equation*}
\begin{aligned}[c]
\cvec{z}_t &= f(\cvec{z}_{t-1},\cvec{a}_{t-1}) + \cvec{\epsilon}_t, \quad \cvec{\epsilon}_t \sim \mathcal{N}(\cvec{0}, \cmat{\Sigma^{trans}}), \\
\cvec{o}_t &= h(\cvec{z}_t) + \cvec{v}_t, \quad \cvec{v}_t \sim \mathcal{N}(\cvec{0}, \cmat{\Sigma^{obs}}).
\end{aligned}
\end{equation*}

Here $\cvec{z}_t \in \mathcal{Z}$, $\cvec{a}_t \in \mathcal{A}$ and $\cvec{o}_t \in \mathcal{O}$ are the latent states, actions, and observations at time t. The vectors $\cvec{\epsilon}_t$ and $\cvec{v}_t$ denote zero-mean Gaussian transition and observation noise respectively.

§.§ Linear Gaussian SSMs (Kalman Filters).

For SSMs defined in <ref>, when $f$ and $h$ are linear/locally linear, inference can be performed efficiently via exact inference. Such SSMs are referred to as Kalman Filters, which were developed independently of PGM Literature. The Kalman filter kalman1961new works by iteratively answering two queries, (i) estimating the prior marginal (predict step) and (ii) updating the posterior given current observations (update step), which are detailed below:

Predict Step
During the prediction step the transition model $\cmat{A}$ is used to infer the current prior state estimate $p(\cvec{z}_t|\cvec{w}_{1:t-1}) = \mathcal{N}(\cvec{z}_t^-,\cmat{\Sigma}_t^-)$, i.e., a-priori to the observation, from the previous posterior estimate $p(\cvec{z}_{t-1}|\cvec{w}_{1:t-1}) = \mathcal{N}\left(\cvec{z}_{t-1}^+, \cmat{\Sigma}_{t-1}^+ \right)$,  by

\begin{equation}
\begin{aligned}
\cvec{z}^-_{t} &= \cmat{A} \cvec{z}^+_{t-1} \\ \text{and} \quad \cmat{\Sigma}^-_{t} &= \cmat{A} \cmat{\Sigma}_{t-1}^+ \cmat{A}^T + \cmat{\Sigma}^\mathrm{trans}.
\end{aligned}
\end{equation}

Update Step
The prior estimate is then updated using the current observation $\cvec{w_t}$ and the observation model $\cmat{H}$ to obtain the posterior estimate $\left(\cvec{z}_t^+, \cmat{\Sigma}_t^+ \right)$, i.e.,

\begin{equation}
\begin{aligned}
\cvec{z}_t^+ &= \cvec{z}_t^- + \cmat{Q}_t \left(\cvec{w}_t - \cmat{H} \cvec{z}_t^-  \right) \\ \text{and} \quad   \cmat{\Sigma}^+_t &= \left(\cmat{I} - \cmat{Q}_t \cmat{H} \right) \cmat{\Sigma}^-_t \\ \text{, with} \quad \cmat{Q}_t &= \cmat{\Sigma}^-_t \cmat{H}^T\left(\cmat{H} \cmat{\Sigma}^-_t \cmat{H}^T  + \cmat{\Sigma}^\mathrm{obs} \right)^{-1},
\end{aligned}
\label{eq:KalmUp}
\end{equation}

where $\cmat{I}$ denotes the identity matrix.
The matrix $\cmat{Q}_t$ is referred to as the Kalman gain.
The whole update step can be interpreted as a weighted average between state and observation estimate, where the weighting, i.e., $\cmat{Q}_t$, depends on the uncertainty about those estimates.

§.§ Deep Kalman Filters

To model dynamics under partially observable scenarios, state space models, in particular the Kalman filter kalman1960new, have been recently integrated with deep learning, by modelling dynamics in learned latent space. We discuss two such works haarnoja2016backprop,becker2019recurrent, that borrow concepts from deep learning and graphical model communities, where the architecture of the network is informed by the structure of the probabilistic state estimator. The various exact inference schemes in the learnt latent spaces derived and executed in the thesis are motivated by this work.

Backprop Kalman Filter
The BackpropKF haarnoja2016backprop uses a deep convolutional encoder to encode a high dimensional observation $\cvec o$ to a latent observation $\cvec w_{t} = \textrm{enc}_w(\cvec o_{t})$. In addition to the latent observation, the encoder also learns to output the uncertainty of this observation, i.e., $\cvec \sigma_{o_t} = \textrm{enc}_{\sigma}(\cvec o_{t}).$ Due to the non-linear observation encoder, a simplified linear Gaussian state space model / (extended) Kalman filter with a known transition model is used in the latent space for state estimation. The parameters of the latent state
distribution are directly optimized as a deterministic computation graph via backpropagation through time. This scheme can also be motivated using Koopman theory koopman1931hamiltonian,mondal2023efficient.

Recurrent Kalman Networks (RKN)

Recurrent Kalman Network (RKN) from <cit.>. It embeds Kalman filter update equations into the latent space of a deep encoder-decoder network. The Kalman predict step gives the current latent prior $\left(\mathbf{z}_t^{-}, \boldsymbol{\Sigma}_t^{-}\right)$using the last posterior $\left(\mathbf{z}_{t-1}^{+}, \boldsymbol{\Sigma}_{t-1}^{+}\right)$and subsequently update the prior using the latent observation $\left(\mathbf{w}_t, \boldsymbol{\sigma}_t^{\mathrm{obs}}\right)$. The factorized representation as shown in Figure <ref> of $\Sigma_t$ converts matrix inversions to scalar operations. Further it allows splitting the latent state $\mathbf{z}_t$ to the observable units $\mathbf{p}_t$ as well as the corresponding memory units $\mathbf{m}_t$ as discussed in Section <ref>. Finally, a decoder is tasked to reconstruct the sensory information.

Recurrent Kalman Network (RKN) becker2019recurrent, builds upon this idea but learns a locally linear transition model as opposed to a known model as in BackpropKF. Most importantly, the authors come up with clever factorized latent state representation as shown in Figure <ref>, that avoids expensive matrix inversions in the Kalman Update Step. We discuss two important implications of this factorized representation below:

The latent state beliefs (mean $\mu$ and covariance $\Sigma$) in <cit.> is factorized into an upper (observed) and lower (unobserved) part as shown in the figure. The lower part can be dubbed as "memory states" and is assumed to keep a memory of information accumulated over time, because of the specific structure of the observation model. The side covariances ($\sigma^s$) are learnt to capture the correlation between the two halves.

* Matrix inversions as scalar operations: Kalman observation update, which is an instance of Bayesian inversion via conditioning involves high dimensional matrix inversions that are expensive to evaluate and hard to backpropagate for end-to-end learning. Hence, <cit.> introduce a factorization of the belief $p(\cvec z_t|\cvec o_{1:t}, \cvec a_{1:t-1}) = \mathcal{N}(\cvec z_t^+, \cvec \Sigma_t^+)$ such that only the diagonal and one off-diagonal vector of the covariance need to be computed, i.e.

$$\cvec \Sigma_t^+ = \left[\begin{array}{cc}\cvec \Sigma_t^{u,+} & \cvec \Sigma^{s,+}_t \\ \cvec \Sigma^{s,+}_t & \cvec \Sigma^{l,+}_t  \end{array} \right], \textrm{ with } \cvec \Sigma_u^+ = \textrm{diag}(\cvec \sigma_t^{u,+}), \; \cvec \Sigma_l^+ = \textrm{diag}(\cvec \sigma_t^{l,+}) $$

and $\cvec \Sigma_{s}^+ = \textrm{diag}(\cvec \sigma_{t}^{s,+})$. Using this factorization, the Kalman filter update step described in Equation <ref> can be executed using only scalar divisions, which are computationally efficient and simpler to implement in back-propagation. Consequently, the closed-form equations for updating the mean can be succinctly expressed using the following scalar equations:

\begin{align*}
\boldsymbol{z}_t^+ = \boldsymbol{z}_t^- +
\left[\begin{array}{c} \boldsymbol{\sigma}^\mathrm{u,-}_t \\
\boldsymbol{\sigma}^\mathrm{l,-}_t \end{array}\right]
\odot
\left[\begin{array}{c}\boldsymbol{w}_t - \boldsymbol{z}^{\mathrm{u},-}_t \\
\boldsymbol{w}_t - \boldsymbol{z}^{\mathrm{u},-}_t  \end{array}\right]
\oslash
\left[\begin{array}{c}  \boldsymbol{\sigma}_t^{\mathrm{u},-} + \boldsymbol{\sigma}_t^\mathrm{obs}
\\  \boldsymbol{\sigma}_t^{\mathrm{u},-} + \boldsymbol{\sigma}_t^\mathrm{obs} \end{array}\right],
\end{align*}

The corresponding equations for the variance update can be expressed as the following scalar operations,

\begin{align*}
\boldsymbol{\sigma}^{\mathrm{u},+}_t &= \boldsymbol{\sigma}^{\mathrm{u},-}_t \odot \boldsymbol{\sigma}^{\mathrm{u},-}_t \oslash \left( \boldsymbol{\sigma}_t^{\mathrm{u},-} + \boldsymbol{\sigma}_t^\mathrm{obs} \right),  \\
\boldsymbol{\sigma}^{\mathrm{s},+}_t &= \boldsymbol{\sigma}^{\mathrm{u},-}_t \odot \boldsymbol{\sigma}^{\mathrm{s},-}_t \oslash \left( \boldsymbol{\sigma}_t^{\mathrm{u},-} + \boldsymbol{\sigma}_t^\mathrm{obs} \right), \\
\boldsymbol{\sigma}^{\mathrm{l},+}_t &= \boldsymbol{\sigma}^{\mathrm{l}, -}_t - \boldsymbol{\sigma}^{\mathrm{s},-}_t \odot \boldsymbol{\sigma}^{\mathrm{s},-}_t \oslash \left( \boldsymbol{\sigma}_t^{\mathrm{u},-} + \boldsymbol{\sigma}_t^\mathrm{obs} \right),
\end{align*}

where $\odot$ denotes the elementwise vector product and  $\oslash$ denotes an elementwise vector division.
These factorization assumptions form the basis for the derivations of exact inference routines for all proposed models in this dissertation.

* Effective Estimation with Correlated Memory and Observation in Latent Variables: Choosing the observation model as $\cvec H = [\cvec I,  \cvec 0]$ allows for the division of the latent state vector into two distinct components. The latent state $\cvec z_t = [\cvec p_t^T, \cvec d_t^T]^T$ has twice the dimensionality of the latent observation $\cvec w_t$ and only the first half of the latent state, i.e., $\cvec p_t$, can be observed. The upper part utilizes the identity matrix in $\cvec H = [\cvec I,  \cvec 0]$ to directly extract information from the observations. Meanwhile, the second lower part remains unobservable and is meant to hold information inferred over time, such as velocities in ordinary dynamical systems or images. The key aspect that contributes to the effectiveness of this choice is the selection of the covariance matrix structure discussed previously. The covariance matrix is designed to incorporate both diagonal and off-diagonal elements, ensuring that the correlation between the memory and the observation parts is effectively learned in the off-diagonal part. On the contrary, if we were to utilize a pure diagonal covariance structure, it would not update the memory units (the later half) or their variance adequately during the Observation/Task/Kalman Update step. Thus, the second half of the latent state, i.e., $\cvec d_t$, serves as derivative or velocity units that can be used by the model to estimate the change of the observable part of the latent state.

In this thesis, we rely on the same factorization assumption for inference in all proposed formalisms to account for "memory" accumulated over time.

§ DEEP BAYESIAN AGGREGATION.

To aggregate information from a set of observations into a consistent representation, <cit.> introduce Bayesian aggregation in the context of Meta-Learning and Neural Process garnelo2018neural literature. Since the aggregation scheme was introduced in the context of multitask latent variable models, we call this latent variable $l$ the task variable in this section. The generative model for the aggregation of observations $n$ is shown in Figure <ref>. <cit.> derive a closed-form update

Generative model for the latent task variable ($l$) inference. The hollow arrows are deterministic transformations leading to implicit distribution $r_{n}$ using an action set encoder.

rule for the posterior $p(\cvec l| \cvec r_{1:N})$ using the Bayes rule, given an observation model of the form,
$$p(\cvec{r}_{n}|\cvec{l}) = \mathcal{N}\left(\cvec{r}_n|\cvec H \cvec{l},\textrm{diag}\left(\cvec{\sigma}_{n}\right)\right)$$
with $\cvec H = \cvec I$ and a prior $p(\cvec l) = \mathcal{N}(\cvec \mu_0, \textrm{diag}(\cvec \sigma_0))$. The Gaussian assumption allows us to obtain a closed-form solution for the posterior estimate of the latent task variable $p(\cvec{l}|\mathcal{\cmat{C}}_{\cvec{l}})$, based on Gaussian conditioning. The factorization assumption further simplifies this update rule by avoiding computationally expensive matrix inversions into a simpler update rule, as follows,

\begin{equation}
\begin{aligned}
\cvec{\sigma_{l}} = \left( \left(\cvec{\sigma}_{0}\right)^\ominus+ \sum_{n=1}^N \left(\cvec{\sigma}_n\right)^\ominus\right)^\ominus, \quad \cvec{\mu}_{\cvec{l}} = \cvec{\mu}_{0} + \cvec{\sigma}_{\cvec{l}} \odot \sum_{n=1}^N \left(\cvec{r}_n - \cvec{\mu}_{0}\right) \oslash \cvec{\sigma}_n.
\end{aligned}
\label{eq:baeq}
\end{equation}

Here $\ominus$, $\odot$ and $\oslash$ denote element-wise inversion, product, and division, respectively.

Using Deep Set Encoders To Learn Observation Uncertainties Note that the posterior calculation is based on the assumption that we have access to information on the latent (compact) observation $\cvec r_n$ and its corresponding uncertainty $\cvec \sigma_{o_n}$. We learn these using a Deep Set encoder as follows,

\begin{align}
p(\cvec{r}_n|\cvec{l}) = \mathcal{N}\left(\cvec{r}_n|\cvec{l},\textrm{diag}\left(\cvec{\sigma}_{n}\right)\right), \quad \cvec{r}_n = \textrm{enc}_{\cvec{r}}\left(\cvec{x}_n\right), \quad \cvec{\sigma}_n = \textrm{enc}_{\cvec{\sigma}}\left(\cvec{x}_n\right).
\end{align}

The architecture of the Bayesian Aggregation Scheme is shown in <ref>.

Given a set of N observations, the deep set encoder emits a latent representation for each of the observations and their corresponding uncertainty. The set of latent representations are then aggregated via Bayesian aggregation using update rules <ref> to obtain a posterior over the abstract latent task variable $l$, $p(\cvec{l}|\mathcal{\cmat{o}}_{1:N})$. The architecture of the deep set encoder can vary depending on the type of sensory signals aggregated.

Bayesian Aggregation As Probabilistic Attention Intuitively the mean of the latent task variable $\cvec{\mu}_{\cvec{l}}$ is a uncertainty weighted sum of the individual latent observations $\cvec{r}_n$, while the variance of the latent task variable $\cvec{\sigma}^2_{\cvec{l}}$ gives the uncertainty of the abstract representation of the task.  The derived aggregation rules can be thought of as simple "probabilistic attention," where the attention weights are given by the learned uncertainty about the individual sensory observations in the set encoder. This is in line with ideas from computational neuroscience that treat cognitive attention to different sensory signals as a form of "precision weighting" hohwy2013predictive,seth2014cybernetic. The derived update equations have only a linear computational complexity of $O(N)$, when similar deep set operations (self-attention) in transformers <cit.> have a complexity of $O(N^2)$.

Note that computing this posterior is a simplified case of the Kalman update rule used in Gaussian SSMs becker2019recurrent, with no memory units, $\cvec H
= \cvec I$ and no dynamics.
In this thesis, we explore the concept of Bayesian Aggregation to synthesize temporally abstract representations within generative world models. We also derive, in Chapter <ref>, generic closed-form update equations for Bayesian Aggregation that avoid the restrictive assumptions of observation models being identity matrices, as seen in <cit.>.

CHAPTER: RELATED WORKS

In this section, I explore the existing literature on internal-world models from two primary fields: computational neuroscience and machine learning. This dissertation introduces scalable, end-to-end deep learning-based approaches for generative internal-world models, along with their respective inference and learning methods. These approaches aim to implement versions of various internal world models that computational neuroscientists have studied, adapting them within the framework of modern deep learning technologies. The discussions on computational neuroscience also offer insight into why certain representations were chosen for the Bayesian networks outlined in this dissertation from a neuroscientific perspective.

§ COMPUTATIONAL NEUROSCIENCE LITERATURE

§.§ The Predictive Bayesian Brain

An increasingly popular theory in cognitive science claims that the brains are essentially prediction machines hohwy2013predictive. The theory is commonly known as the Bayesian brain knill2004bayesian,polydoros2015real, predictive processing clark2013whatever, and predictive mind hohwy2013predictive, among others. At its most fundamental, predictive brain theory says that perception is the result of the brain inferring the most likely causes of its sensory inputs by minimizing the difference between actual sensory signals and the signals expected on the basis of continuously updated predictive models. Arguably, predictive processing provides the most complete framework to date for explaining perception, cognition, and action in terms of fundamental theoretical principles and neurocognitive architectures seth2014cybernetic,jiang2021predictive.

Helmholtz's Unconscious Inference

Shows a basic perceptual inference problem of figuring out what caused a sensory signal (sound e). The mechanism is analogous to the situation for the brainhohwy2013predictive,seth2014cybernetic. Given that the prior probability of hypothesis (cause) is $h_i: P\left(h_i\right)$ and the likelihood that the evidence e would occur, given $h_i$ is $P\left(e \mid h_i\right)$, the brain computes the posterior probability of hypothesis $h_i$, given the evidence $e: P\left(h_i \mid e\right)$ using internal generative models. The figure shows simplified version of Bayes' rule that puts it together: $P\left(h_i \mid e\right)=P\left(e \mid h_i\right) P\left(h_i\right)$.

Bayesian Brain Hypothesis has its origins in Hermann von Helmholtz's notion of unconscious inference helmholtz1948concerning, proposing that perception is a result of inferential processes by the brain, rather than direct imprinting of the external world onto our sensory apparatus. The observer is typically not aware of such prior assumptions but rather, they are incorporated by the neural circuits subconsciously to compute beliefs over hidden causes through the dynamics of neural activities (thereby implementing perception as “unconscious inference”). As part of the internal model, such priors can be expected to be adapted to the environment that the organism lives in. This idea laid the groundwork for the Bayesian brain hypothesis, which posits that the brain computes probabilities to infer the state of the world from ambiguous sensory inputs and prior beliefs, effectively engaging in a form of statistical reasoning or Bayesian inference.

Bayesian Brain Hypothesis
The Bayesian brain hypothesis knill2004bayesian uses Bayesian probability
theory to formulate perception as a constructive process based on internal or generative models. The underlying idea is that the brain has a model of the world von2013treatise,mackay1956epistemological,neisser2014cognitive,gregory1968perceptual that it tries to optimize using sensory inputs ballard1983parallel,kawato1993forward,kersten2004object,lee2003hierarchical,friston2005theory. This idea is related to analysis by synthesis neisser2014cognitive and epistemological automata mackay1956epistemological. In this view, the brain is an inference machine that
actively predicts and explains its sensationsvon2013treatise,gregory1980perceptions,dayan1995helmholtz. Central to this hypothesis is a probabilistic model that can generate predictions against which sensory samples are tested to update beliefs about their causes. This generative model is decomposed into a likelihood (the probability of sensory data given their causes) and a
prior (the a priori probability of those causes). Perception then becomes the process of inverting the likelihood model (mapping from causes to sensations) to access the posterior probability of the causes, given sensory data (mapping from sensations to causes).

§.§ Kalman Filters As Internal World Models / Spatiotemporal Predictive Coding

The world is dynamic; most of the time, animals receive time-varying stimuli either due to their own movement or due to other moving objects in the environment. This makes the ability to predict future stimuli essential for survival (e.g. predicting the location of predators). Rao's seminal work rao1999optimal from the computational neuroscience literature introduces a foundational mathematical framework that underpins how organisms perceive their environment in a dynamic setting as opposed to static settings discussed previously. This framework posits that visual perception is a stochastic, dynamic process and the task of perception is one of optimally estimating (in a Bayesian sense) the causes of visual events and on a longer time scale, learning efﬁcient spatio-temporal internal models of the visual environment. The spatiotemporal internal generative model is elegantly represented as a stochastic linear dynamical system similar to a Kalman Filter with transition and observation matrices. By employing inference techniques rooted in Kalman filter update equations, Rao's approach provides a Bayesian optimal method for estimating the internal states that govern the dynamics of the mental world model.

Figure from <cit.> presents the challenge organisms face in perceiving the external world through internal models. Without direct access to the world's hidden internal states, organisms rely on sensory measurements to estimate these states, solving the 'inverse' problem to interpret and understand their environment. Though the paper focused on visual systems, the authors argue that the concept extends beyond visual and auditory senses to include internal models for motor systems, emphasizing the broader application of understanding and interacting with the external world through internal representations.

Combining Top-Down Predictions With Bottom-Up Sensory Signals Using Kalman Filter Updates
The proposed Kalman-filter model in <cit.> provides a mathematical framework that demonstrates how top-down expectations (informed by the internal generative model) interact with bottom-up signals (the current sensory input) to produce a robust estimation of the visual environment. The Kalman Filter combines these two sources of information, weighted according to their reliability (inverse variances or “precisions”), to compute the posterior beliefs about the world in a Bayesian optimal manner, which is interpreted as a form of attention. The proposed model allows for the possibility that the organism or agent might want to perform internal simulations of the dynamics of the external world (e.g., for planning) by predicting how future states evolve given a starting state (and possibly actions).

Learning Of Internal Models The work also derived update rules for learning the transition and precision matrices online based on incoming input data, as opposed to hand-coded dynamics models.

<cit.> further provides experimental results that support the model's capability for learning dynamic and static visual stimuli, recognizing objects under various conditions, and segmenting scenes into constituent elements. The work from <cit.>, also referred to as "spatio-temporal predictive coding" in literature, thus sets a foundational stone for exploring the synergies between computational neuroscience and the development of popular machine learning architectures like Deep Kalman Models.

The neuroscientific insights discussed above are exploited in Chapter <ref> of the thesis using modern deep learning techniques.

§.§ Adaptive Behavior In Brain With Multiple Internal Models

One of the most notable abilities of biological creatures is their capacity to adapt their behavior to different contexts and environments (i.e., cognitive
flexibility) through learning. People can learn to call on various responses depending on the situation—for example, independently move the right and left hands when playing an instrument and speak several different languages. Such multitasking abilities are particularly crucial in communication with several people, who each demand subtly different forms of interaction taborsky2012social,parkinson2015repurposed and is a key component of exhibiting social intelligence. The work from <cit.> attempts to understand how the brain attains cognitive flexibility in different contexts by entertaining the possibility of distinct generative models in a context-sensitive setting. <cit.> further present the results of perceptual learning and inference using this form of model selection or structure learning, predicated on an ensemble of generative models. [24]r0.35

The perceptual hierarchy hohwy2013predictive: Processing of causal regularities at different time scales influence each other in a bottom-up–top-down fashion. Sensory input (dark grey arrows pointing up) is met with prior expectations (black arrows pointing down) and perceptual inference is determined in multiple layers of the hierarchy simultaneously, building up a structured representation of the world. This figure simplifies greatly as there are of course not just three well-defined hierarchical levels in the brain.

Using this setup, it is shown that Bayesian model averaging provides a plausible account of how multiple hypotheses can be combined to predict the sensorium, while Bayesian model selection enables perceptual categorization and selective learning.

The generative model proposed in Chapter <ref> of the dissertation have a similar motivation of using a "set of SSMs/world models" for adaptive behaviour to changing tasks.

§.§ Hierarchical Generative Models In The Brain

Researchers in computational neuroscience argue that the previously discussed internal predictive world models are hierarchical in nature friston2008hierarchical,hohwy2013predictive,seth2014cybernetic. The hierarchical nature adeptly mirrors the complex and layered structure of the world itself, capturing the essence of how causes and effects interrelate across different spatial and temporal scales hohwy2013predictive.
These hierarchical generative models overturn classical notions of perception that describe a largely 'bottom-up' process of evidence accumulation or feature detection. Instead, predictive processing proposes that perceptual content is determined by top-down predictive signals emerging from multilayered and hierarchically organized generative models of the causes of sensory signals lee2003hierarchical. These models are continually refined by mismatches (prediction errors) between predicted signals and actual signals across hierarchical levels, which iteratively update predictive models via approximations to Bayesian inference. This means that the brain can induce accurate generative models of hidden environmental causes by operating only on signals to which it has direct access: predictions and prediction errors. It also means that even low-level perceptual content is determined via cascades of predictions flowing from very general abstract expectations, which constrain successively more fine-grained predictions.
The hierarchical generative model proposed in Chapter <ref> takes inspiration from these insights.

§ MACHINE LEARNING LITERATURE

We give a brief overview of notable related works in learning computational mental models of the world in the field of machine learning. A detailed discussion on several of these are provided in Sections <ref>, <ref> and <ref>.

World Models Based On Gaussain Processes

Bayesian nonparametric models, such as Gaussian processes (GPs), are often the dynamics model of choice in Model-Based RL, especially in low-dimensional problems where data efficiency is critical due their robustness in handling uncertainty kocijan2004gaussian,ko2007gaussian,nguyen2008local,grancharova2008explicit,deisenroth2013gaussian,kamthe2018data.

PGM representation adapted from <cit.> of a GP regression for learning dynamics. Squares represent observed variables, and circles represent unknowns. The thick horizontal bar represents a set of fully connected nodes and indicates that the forward dynamic function $f_i$ belongs to a Gaussian field. During posterior inference $f_i$'s act as random variables and are integrated out, which means that every prediction $y_*$ depends on all the other inputs and observations (including the training data) making GPs often computationally challenging.

However, such models introduce additional assumptions on the system, such as the smoothness assumption inherent in GPs with squared-exponential kernels rasmussen2003gaussian. Also, these methods do not scale to high-dimensional environments. Furthermore, these models are inherently single-time scale and are not widely used in tasks where long-horizon predictions and planning are required.

World Models Based On Neural Networks
The advancements in parametric function approximators, such as neural networks (NNs), brought about by the advent of deep learning, have led to the development of a series of world models employing multi-layer perceptrons. These models are cited in several works, including <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, and <cit.>. In contrast to Gaussian processes, these deep learning-based models offer constant-time inference and manageable training with large datasets. They also possess the capability to model more intricate functions, including the non-smooth dynamics frequently encountered in robotics fu2016one,nagabandi2018neural.[20]r0.4

(Top) A dynamics model based on deep neural networks. (Bottom) A graphical representation using PGM of the same model, where the shaded nodes indicate observed variables. These models operate under the assumption that the current states of the internal model of the world are already known and don't need to be inferred. This assumption is limiting, particularly when these models need to be derived from sensory information that is both high-dimensional and fraught with noise.

However, the majority of these studies utilizing NNs have concentrated on deterministic models, which do not align well with the Bayesian perspective of best guesses discussed earlier. <cit.> proposed utilizing an ensemble of neural network models to learn probabilistic dynamics, capturing both the aleatoric uncertainty inherent in the environment and the epistemic uncertainty of the data. This approach allows for more data-efficient control and planning algorithms. Nonetheless, irrespective of being deterministic or stochastic, these NN-based dynamics models face challenges in learning with high-dimensional sensory data or in partially observable domains due to the presumption of known internal world states. Additionally, for long-term predictions, these auto-regressive models are less effective because of the accumulation of prediction errors inherent in autoregressive processes. They also presume static world dynamics and operate at a singular, fine-grained temporal scale, thereby lacking in their ability to serve as a general-purpose world model capable of addressing partial observability, non-stationarity, and hierarchical predictions across multiple time scales.

World Models Based On State Space Models
Classical State-Space Models (SSMs) are well-regarded for their clear inference processes and interpretable outcomes, but they often struggle with scalability and efficiency when applied to high-dimensional data or large datasets. To overcome these limitations, recent advances have introduced deep SSMs, which integrate neural networks to provide efficient and scalable inference for complex datasets. The works of <cit.> have leveraged neural network architectures to derive closed-form solutions for the forward inference algorithms in SSMs, leading to state-of-the-art performance in tasks like state estimation and dynamics prediction. Other studies, such as those of <cit.>, have focused on applying variational approximations to facilitate learning and inference within SSMs.

Despite these advances, many recurrent state-space models are based on the assumption of fixed dynamics, which does not align with the changing conditions observed in real-world scenarios, such as in robotics. To address the limitations of fixed dynamics, research by <cit.> has introduced models that incorporate additional discrete switching latent variables to account for changing or multi-modal dynamics. However, these models typically operate on a single time scale, with both continuous and discrete levels operating at a fine-grained temporal resolution. They also do not adhere to a strictly top-down approach, as the top-layer latent variables depend on the predictions from the lower layer, rendering these models not suitable for long-horizon predictions and reasoning. Moreover, the introduction of discrete states complicates the learning and inference processes.

CHAPTER: ACTION CONDITIONAL DEEP SSMS: PROBABILISTIC WORLD MODELS ON A SINGLE TIME SCALE

This chapter is based on "Action Conditional Recurrent Kalman Networks for Forward and Inverse Dynamics Learning" shaj2020action.

In our quest to formalize a computational mental model of the world, we start with a simple and widely used Bayesian network called the State Space Models (SSMs). Though it operates at a single time scale (often the time scale at which observations are recorded), and hence may not satisfy the expressiveness and hierarchical structure of the mental models that the Neuroscience community advocates (Section <ref>), SSMs have several computational advantages. The chain structure of SSMs allows for tractable exact inference via message passing allowing for closed-form update rules under linear and Gaussian assumptions. When these updates are performed in the learnt latent space of a Deep Network, these linear assumptions are no longer limiting. These deep models learn a latent space where these linear assumptions work in practice. The Deep SSMs with exact inference are known in the literature as Deep Kalman models and we use one such model namely Recurrent Kalman Networks (RKN) becker2019recurrent discussed in the preliminaries Section <ref> as our foundation for learning a world model.

Research Objective Deep Kalman Models including RKN were primarily designed for state estimation from sensory observations in simple unactuated dynamical systems. In the real world, an agent's actions actively shape its environment, and ignoring this causal relationship in the mental models of the world can lead to inaccurate predictions and suboptimal decision-making. Hence we try to answer these questions in this chapter of the thesis:

* Can deep Kalman models be adapted for the task of world modelling?

* How to incorporate control action conditioning in the latent dynamics in a principled manner such that causal relations are respected?

* How good are these models in modelling the dynamics of complex robotic systems when faced with non-markovian, partially observable and non-stationary settings?

§ RELATED WORKS

Action Conditional Predictive Models. In model-based control and reinforcement learning (RL) problems, learning to predict future states and observations conditioned on actions is a key component.
Approaches such as <cit.>, <cit.>, <cit.>, <cit.>, <cit.> try to achieve this by using traditional models like GPs, feed forward neural networks, or LSTMs.
The action conditioning is realized by concatenating the input observations and action signals nagabandi2018neural,deisenroth2011pilco or via factored conditional units taylor2009factored, where features from some number of inputs modulate multiplicatively and are then weighted to form network outputs.
In each of these approaches action conditioning happens outside of the recurrent neural network cell which leads to sub-optimal performance as observations and actions are treated similarly.

Predictive Models Of Robotic Agents. Due to the increasing complexity of robot systems, analytical models are more difficult to obtain. This problem leads to a variety of model estimation techniques which allow the roboticist to acquire models from data.
When learning the model, mostly standard regression techniques with Markov assumptions on the states and actions are applied to fit either the forward or inverse dynamics model to the training data.
Authors used Linear Regression schaal2002scalable,haruno2001mosaic, Gaussian Mixture Regression calinon2010learning,khansari2011learning, Gaussian Process Regression nguyen2009model,nguyen2010using, Support Vector Regression ferreira2007simulation, and feed forward neural networks polydoros2015real.
However the Markov assumptions are violated in many cases, e.g., for learning the dynamics of complex robots like hydraulically and pneumatically actuated robots or soft robots.
Here, the Markov assumptions for states and actions no longer hold due to hysteresis effects and unobserved dependencies.
Similarly, learning the dynamics of robots in frictional contacts with unknown objects also violates the Markov assumption and requires reliance on recurrent models which can take into account the temporal dynamic behavior of input sequences while making predictions.
Recently, using RNNs such as LSTMs or GRUs has been attempted.
Those scale easily to big data, available due to the high data frequencies, and allow learning in $O(n)$ time rueckert2017learning.
However, the lack of a principled probabilistic modelling and action conditioning makes them less reliable to be applied for robotic control applications.

§ ACTION CONDITIONAL RECURRENT KALMAN NETWORKS

We extend the Recurrent Kalman Network approach to the task of world modelling / forward dynamics learning with the following innovations: (i) We use an action-conditional prediction update which provides a natural way to incorporate control inputs into the latent dynamical system.  (ii) We modify the architecture to perform prediction instead of filtering and learn accurate forward dynamics models for multiple robots with complex actuator dynamics. We refer to the resulting approach as action-conditional RKN (ac-RKN).

§.§ Action Conditioning

Graphical models for actions (green). (Left) Treated as fake observations as in <cit.>. (Right) ac-RKN treating actions in a principled manner by capturing its causal effect on the state transition via additive interactions with the latent state. The hollow arrows denote deterministic transformation leading to implicit distributions.

To achieve action conditioning within the recurrent cell, we modify the transition model proposed in <cit.> to include  a control model $\vec{b}(\vec a_t)$ in addition to the locally linear transition model $\vec{A}_t$.
The control model $\vec{b}(\vec a_t)$ is added to the locally
linear state transition, i.e., $\vec{z}^-_{t+1} = \vec{A}_t \vec{z}^+_{t} + \vec b(\vec a_t)$.
As illustrated in <ref>, this formulation of latent dynamics captures the causal effect of the actions variables $\vec{a}_t$ on the state transitions in a more principled manner than including the action as an observation, which is the common approach for action conditioning in RNNs becker2019recurrent.
The control model $\vec{b}(\vec a_t)$ can be represented in several ways, i.e.:

(i) Linear: $\vec{b}_{l}(\vec a_t) =\vec{B} \vec a_t$, where $\vec{B} $ is a linear transformation matrix.

(ii) Locally-Linear: $\vec{b}_{m}(\vec a_t) = \vec{B}_t \vec a_t$, where $\vec{B}_t = \sum_{k=0}^K \beta^{(k)}(\vec{z_t}) \vec{B}^{(k)}$ is a linear combination of k linear control models $\vec{B}^{(k)}$. A small neural network with softmax output is used to learn $\beta^{(k)}$.

(iii) Non-Linear: $\vec{b}_n(\vec a_t) = \vec{f}(\vec a_t)$, where $\vec{f}(.)$ can be any non-linear function approximator. We use a multilayer neural network regressor with ReLU activations.

Note that the prediction step for the variance is not affected by this choice of action conditioning, i.e., $\vec{\Sigma}^-_{t+1} = \vec{A}_t \vec{\Sigma}_{t}^+ \vec{A}_t^T + \vec{I} \cdot \vec{\sigma^\mathrm{trans}}$ as the action is known and not uncertain.
Thus, unlike the state transition model $\vec{A}_t$, we do not need to constrain the model $ \vec{b}$ to be linear or locally linear, as it neither affects the Kalman gain nor how the covariances are updated.

Hence, we use the nonlinear approach, $\vec{b}_n(\vec a_t)$, as it provides the most flexibility and achieves the best performance, as shown in Section <ref>.
The Kalman update step is also unaffected by the new action-conditioned prediction step and therefore remains as presented in <cit.>.

Our principled treatment of control signals is crucial for learning accurate forward dynamics models, as detailed in Section <ref>. Moreover, the disentangled representation of actions in the latent space gives us more flexibility in manipulating the control actions for different applications including inverse dynamics learning (Appendix <ref>), extensions to model-based reinforcement learning, and planning. For long-term prediction using different action sequences, we can still apply the latent transition dynamics without using observations for future time steps (skipping the observation update step).

§.§ Ac-RKN as Single Time Scale Probabilistic World Models

Using internal world models to envision the future: With past and present observations ($w_{1:2}$) as a basis, Ac-RKN projects future world states ($z_{3:6})$ by considering a sequence of action signals ($a_{1:5}$).

World models attempt to learn a compact and expressive representation of the environment dynamics from observed data. These models can predict possible future world states and associated uncertainties as a function of an imagined action sequence. More formally, we want to learn a forward dynamics model of the world $f: \vec o_{1:t}, \vec a_{1:t} \mapsto \vec o_{t+1}$ that predicts the next observation $\vec o_{t+1}$ given the histories of observation $\vec o_{1:t}$ and actions $\vec a_{1:t}$.

Formulating action-conditional future predictions as inference tasks within a Gaussian State Space Model (SSM) in deep latent spaces enables the development of probabilistic world models that exhibit the following properties.

* The propagation of uncertainties ($\cvec \Sigma_t$) about the belief of our future world states in addition to the mean estimates ($\cvec{z}_t$).

* Learning under partial observability and missing observations. The RKN cell is well equipped to handle missing observations, as the update step can just be omitted in this case, i.e. the posterior is equal to the prior if no observation is available. In this case, we repeatedly apply the ac-RKN prediction step while omitting the Kalman update.

* Capturing non-markovian dynamics The model captures non-markovian dynamics in a principled way, since we predict the future belief states $p(\vec z_{t+1}| \vec w_{1:t}, \vec a_{1:t})$ conditioned on all the past observations and the control actions given so far. These beliefs are obtained in closed form using the forward inference algorithm (action conditional Kalman predict steps).

In line with theories from neuroscience discussed in Section <ref>, this results in a generative model that makes Bayesian best guesses about the world states based on the action conditional prior beliefs (top-down predictions) and bottom-up sensory signals/ observations when available. Neuroscience research hohwy2013predictive,seth2014cybernetic,seth2021being similarly posits that the Brain is not only tasked with the challenge of figuring out the most likely causes of its sensory signals but also figuring out how relevant the sensory inputs are. Similarly, the Ac-RKN model learns to estimate the reliability of the incoming sensory information by predicting uncertainties of the sensory signals via the observation encoder (bottom-up processing). For a visual representation of these top-down and bottom-up predictions in the Ac-RKN, please see Figure <ref>.

§.§ End To End Learning Via Backpropagation

The network is tasked to minimize the prediction errors by maximizing the posterior predictive log-likelihood, which is given below for a single trajectory, i.e.,

\begin{align}
\label{eq:objective}
L & =\sum_{t=1}^H \log p(\cvec{o}_{t+1}|\cvec{w}_{1:t}, \cvec{a}_{1:t}) \nonumber \\ & = \sum_{t=1}^H \log \int p(\cvec{o}_{t+1}|\cvec{z}_{t+1})  p(\cvec{z}_{t+1}|\cvec{w}_{1:t}, \cvec a_{1:t}) d\cvec{z}_{t+1} \nonumber \\
\end{align}

The extension to multiple trajectories is straightforward and omitted to keep the notation uncluttered. Here, $\cvec{o}_{t+1}$ is the ground truth observations at the time step $t+1$ which needs to be predicted from all observations up to time step $t$.

Approximating the likelihood We employ a Gaussian approximation of the posterior predictive log-likelihood of the form $ p(\cvec{o}_{t+1}|\cvec{w}_{1:t}, \cvec{a}_{1:t}) \approx \mathcal{N}(\cvec{\mu}_{\cvec{o}_{t+1}},\textrm{diag}(\cvec{\sigma}_{\cvec{o}_{t+1}}))$ where we use the mean of the prior belief $\cvec{\mu}_{z_{t+1}}^-$ to decode the predictive mean, i.e, $\cvec{\mu}_{\cvec{o}_{t+1}} =  \textrm{dec}_{\cvec{\mu}}(\cvec{\mu}_{z_{t+1}}^{-})$ and the variance estimate of the prior belief to decode the observation variance, i.e., $\cvec{\sigma}_{o_{t+1}} = \textrm{dec}_{\sigma}(\cmat{\Sigma}_{z_{t+1}}^{-})$. This approximation can be motivated by a moment-matching perspective and allows for end-to-end optimization of the log-likelihood without using auxiliary objectives such as the ELBO becker2019recurrent. Thus the approximate Gaussian predictive log-likelihood for a single sequence is then computed as

\begin{align}
\label{eq: likeli-acrkn}
&\mathcal{L}\left(\cvec{o}_{(1:T)}\right) = \\&
\dfrac{1}{T} \sum_{t=1}^T  \log \mathcal{N}\left(\cvec{o}_t \bigg| \textrm{dec}_{\mu}(\cvec{z}_t^+), \textrm{dec}_{\Sigma} (\cvec{\sigma}^{\mathrm{u},+}_t,
\cvec{\sigma}^{\mathrm{s},+}_t,
\cvec{\sigma}^{\mathrm{l},+}_t )\right), \nonumber
\end{align}

where $\textrm{dec}_{\mu}(\cdot)$ and $\textrm{dec}_{\Sigma}(\cdot)$ denote the parts of the decoder that are responsible for decoding the latent mean and latent variance respectively.

Schematic diagram of ac-RKN for forward dynamics learning. Action conditioning is implemented by adding a latent control vector $\vec b(\vec a_t)$ to the RKN dynamics model. The output of the prediction stage, which forms the prior for the next time step $\left(\vec{z}_{t+1}^-, \vec{\Sigma}_{t+1}^- \right)$ is decoded to obtain the prediction of the next observation.

Variations Of The Learning Objective for Learning Real Robot Dynamics
In this part of the thesis, we specifically look at the learning challenges associated with modeling the dynamics of real robot systems and propose different learning objectives. We explicitly focus on non-Markovian systems, where the observations are typically joint angles and, if available, joint velocities of all degrees of freedom of the robot.
In our experiments, we aim to predict the joint angles at the next time step given a sequence of joint angles and control actions.
Due to unmodeled effects such as hysteresis or unknown contacts, the state transitions are non-Markovian even if joint angles and velocities are known. We can assume that the observations are almost noise-free, since the measurement errors for our observations (joint angles) are minimal. Nevertheless, as the observations do not contain the full state information of the system, we still have to model uncertainty in our latent state using the RKN.

1. Mean Predictions To ensure a fair comparison with deterministic world models such as LSTM and GRU, we train the model with a variation of the objective in Equaton <ref>, where only the mean of the latent states are decoded, while the variance units are kept fixed, resulting in an RMSE loss.

\begin{align*}
\mathcal{L}_{\textrm{fwd}} = \sqrt{
\dfrac{1}{T}\sum_{t=1}^T  \parallel \left(\vec{o}_{t+1} -  \vec{o}_{t}\right) - \textrm{dec}\left(\vec{z}^-_{t+1}\right) \parallel^2}.
\end{align*}

We also observed that in use cases where only mean predictions are required, this results in slightly better predictions.

2. High-frequency Predictions The function $f$ can be challenging to learn when the observations, i.e. joint positions and velocities, of two subsequent time steps are too similar (for example, our Franka Robot's operating frequency in 1000Hz). In this case, the action has seemingly little effect on the output, and the learned strategy is often to copy the previous state for the next-state prediction. This difficulty becomes more pronounced as the time step between states becomes smaller, e.g., $1$ms, as minor errors in absolute states estimates can already create unrealistic dynamics.  Therefore, a standard method for model learning is to predict the normalized differences between subsequent observations or states instead of the absolute values <cit.>, that is, during training, the next predicted observation is $\hat{\vec o}_{t+1} = \vec o_t + \textrm{dec}\left(\vec{z}_{t+1}^-\right)$, where $\vec{o}_t$ is the true observation in $t$ and $\textrm{dec}\left(\vec{z}_{t+1}^-\right)$ is the output of the actual decoder network. During inference, we introduce an additional memory $\vec{m}_t$ that stores the last prediction and uses it as an observation in case of a missing observation, that is, $\vec{\hat{o}}_{t+1} = \vec{m}_t + \textrm{dec}\left(\vec{z}_{t+1}^-\right)$ where we use the true observation, $\vec{m}_t= \vec{o}_t$ if available, and the predicted observation, $\vec{m}_t= \hat{\vec{o}}_t$, otherwise.

Irrespective of the learning objectives/loss functions used, we still have to model uncertainty in our latent state using the RKN as the observations do not contain the full state information of the system.
The architecture of the ac-RKN is summarized in Figure <ref>.

Note that we decode the prior mean $\vec{z}^-_{t+1}$ to predict $\vec o_{t+1}$.
The prior mean $\vec{z}^-_{t+1}$ integrates all information up to time step $t$, including $\vec a_t$, but already denotes the belief for the next time step.
Hence, as we are interested in prediction, we work with this prior.
On the contrary, <cit.> used posterior belief, as their goal was filtering rather than prediction.

§.§ Self Supervised Training For Multi-Step Prediction

Self-supervised training for multi-step ahead predictions by tasking the model to "fill in" the masked observations.

Using the training objective(s) discussed in Section <ref> results in models that are good in one-time step prediction, but typically perform poorly in long-term predictions as the loss assumes that observations are always available up to time step $t$. To increase the performance of the long-term prediction, we can treat the long-term prediction problem as a case of the problem of “missing value”, where the missing observations occur in future time steps. Thus, to train our model for long-term prediction, we randomly mask a fraction of observations as shown in Figure <ref> and explicitly task the network to impute the missing observations, resulting in a strong self-supervised learning signal for long-term prediction with varying prediction horizon length.

§ EXPERIMENTS

In this section, we discuss the experimental evaluation of ac-RKN on learning forward dynamics models on robots with different actuator dynamics. A full listing of hyperparameters can be found in the Supplementary Material. We compare ac-RKN to both standard deep recurrent neural network baselines (LSTMs, GRUs and standard RKN), analytical baselines based on classical rigid body dynamics, and nonrecurrent baselines like FFNN. For the RKN and LSTM baselines, we replaced the ac-RKN transition layer with generic LSTM and RKN layers. For recurrent models, the parameters of the recurrent cell are tuned through hyperparameter optimization using GPyOpt gpyopt2016, but the size of the encoder and decoder are similar. The observations and actions are concatenated as in <cit.> and <cit.> for all baseline experiments during the forward dynamics learning, that is, we treat actions as extended observations. The data, i.e., the states or observations, actions, and targets, are normalized (zero mean and unit variance) before training. We denormalize the predicted values during inference and evaluate their performance in the test set. We evaluate the model using RMSE to ensure a fair comparison with deterministic models such as FFNNs and LSTMs.

The experiments are performed on data from robots with different actuator dynamics. From left to right, these include: Hydraulically actuated BROKK-40 <cit.>, pneumatically actuated artificial muscles <cit.>, Franka Emika Panda Robotic Arm

We evaluate the ac-RKN on three robotic systems with different actuator dynamics each of which poses unique learning challenges:  (i) Hydraulically Actuated BROKK-40 Robot Arm. The data consists of measured joint positions and the input current to the controller sampled at 100Hz from a hydraulic BROKK-40 demolition robot taylor2013state. The position angle sensors are rotary linear potentiometers. We chose a similar experimental setup and metrics as in <cit.> to ensure a fair comparison. We trained the model to predict the joint position 2 seconds, i.e.  200 time-steps, into the future, given only control inputs. Afterwards, the model receives the next observation and the prediction process repeated. Learning the forward model here is difficult due to inherent hysteresis associated with hydraulic control.
(ii) Pneumatically Actuated Musculoskeletal Robot Arm. This four DoF robotic arm is actuated by Pneumatic Artificial Muscles (PAMs) buchler2016lightweight.
Each DoF is actuated by an antagonistic pair of PAMs, yielding a total of eight actuators.
The robot arm reaches high joint angle accelerations of up to 28,000 $\textrm{deg}/s^2$ while avoiding dangerous joint limits thanks to the antagonistic actuation and limits on the air pressure ranges.
The data consists of trajectories of hitting movements with varying speeds while playing table tennis buchler2020learning and is recorded at 100Hz. The fast motions with high accelerations of this robot are complicated to model due to hysteresis.
(iii) Franka Emika Panda Arm.  We collected the data from a 7 DoF Franka Emika Panda manipulator during free motion at a sampling frequency of 1kHz.
It involved a mix of movements of different velocities from slow to swift motions.
The high frequency, together with the abruptly changing movements results in complex dynamics which are interesting to analyze.

RMSE      NLL
ac-RKN    0.0144     6.247
RKN        0.0282     4.930
LSTM    0.2067      0.952
GRU    0.2015      1.186

Hydraulic Brokk 40

Hydraulic Arm Predictions

(a) Performance comparison of various recurrent models for the BROKK-40 hydraulically actuated robot arm. (b) Visualization of the predicted trajectories of the hydraulic arm for 200 step ahead predictions.

Multi Step Ahead Prediction. Figures <ref> and <ref> summarize the test set performance for each of these robots.  We benchmarked the performance of ac-RKN with state-of-the-art deterministic deep recurrent models (LSTM and GRU) and non-recurrent models like feed-forward neural network (FFNN). In all three robots, the ac-RKN gave much better multi-step ahead prediction performance than these deterministic deep models. The performance improvement was more significant for robots with hydraulic and pneumatic actuator dynamics due to explicit non-markovian dynamics and hysteresis, which is often difficult to model via analytical models. We also validated the performance improvement due to our principled action-conditioning in the latent transition dynamics by comparing it with RKN becker2019recurrent, which treats actions as part of the observations by `concatenation'. In all three robots, our principled treatment brought significant improvement in performance. Figures <ref> and <ref> show the predicted trajectories for the BROKK and the pneumatically actuated robot arm.

Plots show the comparison of different action-conditioning models discussed in Section <ref> for  (left) Franka Emika Panda Robot and (right) Pneumatic Muscular Robot Arm. The plots clearly show that the predictions given by our approach are by far the most accurate.

Comparison with Analytical Model. We were also interested in comparing these with analytical models of Franka. For the analytical model, in addition to the inertia properties of the links, Coulomb friction was also identified for each joint. A detailed description of the same can be found in Appendix C. As seen in Figure <ref>, the performance of the analytical model outperformed ac-RKN for step-by-step predictions, but for multistep forward predictions, data-driven models had a clear advantage with ac-RKN giving the most accurate results.

Impact of control input models

Ablation Study for Action-Conditioning.
We evaluated different models for action conditioning as discussed in Section <ref>. The resulting evaluation can be seen in Figure <ref> and shows the advantage of using nonlinear models for the additive action conditioning of the latent dynamics.
Inverse Dynamics Learning
In addition, we also adapted ac-RKN for inverse dynamics learning tasks. Since this is out of scope of the main topic of the thesis of World modeling, we report the details in Appendix <ref>.

Visualization of the predicted trajectories of the pneumatic arm for 5-step ahead (top) and 10-step ahead (middle) predictions for one of the joints. The bottom plot shows the corresponding pressure (actions) applied in bars. The figure clearly shows that the predictions given by our principled action-conditional scheme are by far the most accurate.

§ CONCLUSION

In this section, we modified a contemporary deep Kalman filter approach for world modeling by incorporating action/control signals into the generative process giving them the capability to predict the sensory consequences of actions. By conditioning the actions on the latent state while maintaining the causal relationships between these entities, we enable the execution of interventions and counterfactual analyses with these signals. Such capabilities are essential for decision-making and planning processes using these models. This design facilitates the learning of precise forward dynamics models for complex robots and scenarios lacking analytical models. We validated the efficacy of our method on robots equipped with hydraulic, pneumatic, and electric actuators in settings with stationary dynamics. Initial experiments indicated that the current generative model, the ac-SSM, falls short in handling non-stationary dynamics. We plan to address this limitation in the subsequent chapter.

CHAPTER: HIDDEN PARAMETER SSMS: PROBABILISTIC WORLD MODELS WITH TASK ABSTRACTIONS

This chapter is based on "Hidden Parameter Recurrent State Space Models For Changing Dynamics Scenarios" shaj2022hidden.

One of the most notable abilities of intelligent biological agents is their ability to adapt their behavior to different contexts and environments (i.e., cognitive flexibility) through learning. <cit.> postulates that the brain maintains a collection of generative internal models of the world, and consequential flexibility arises from contextualization and selection on the basis of higher-order beliefs about the most plausible hypothesis, task, or context in play.

(left) We visualize a set of internal models of the world, each depicted as an individual State-Space Model (SSM) as detailed in Chapter <ref>. These models cater to diverse dynamics and tasks. (right) The proposed HiP-SSM framework enables modeling of multitask dynamics using a singular, overarching model. It achieves this through a hierarchical latent task variable, denoted as $l$ that parametrize the latent dynamics.

In this part of the thesis, we focus on the adaptability aspects of world models. For many real-world control applications, an intelligent agent must learn to solve tasks with similar, but not identical, dynamics. A robot playing table tennis may encounter several bats with different weights or lengths, while an agent manipulating a bottle may encounter bottles with varying amounts of liquid. An agent driving a car may encounter many different cars, each with unique handling characteristics. Humanoids learning to walk may face different terrains with varying slopes or friction coefficients. Any real-world dynamical system might change over time due to multiple reasons, some of which might not be directly observable or understood. For example, in soft robotics small variations in temperature or changes in friction coefficients of the cable drives of a robot can significantly change the dynamics. Similarly, a robot may undergo wear and tear over time, which can change its dynamics. Assuming a global model as discussed in Chapter <ref> that is accurate throughout the entire state space or duration of use is a limiting factor for using such models in real-world applications.

Research Objective Given the limitations of existing literature on recurrent models, particularly in modeling the dynamics of nonstationary situations, and the impracticality of employing separate State Space Models (SSMs) for each task in continuously changing systems, this research attempts to answer the following question:

* Is it feasible to develop a formalism that enables the learning of a unified global dynamics model, applicable across multiple tasks and environments, while incorporating a hierarchical latent task variable for task-specific specialization?

* Can we perform learning and inference in a scalable manner with such a generative model ?

* How do these models compare against traditional Recurrent State Space Models (RSSMs) and other contemporary state-of-the-art models in environments with evolving dynamics?

We thus introduce hidden parameter state-space models (HiP-SSM), which allow capturing the variation in the dynamics of different instances through a set of hidden task parameters. We formalize the HiP-SSM and show how to perform inference in this graphical model. Under Gaussian assumptions, we obtain a probabilistically principled yet scalable approach. We name the resulting probabilistic recurrent neural network as Hidden Parameter Recurrent State Space Model (HiP-RSSM). HiP-RSSM achieves state-of-the-art performance for several systems whose dynamics change over time. Interestingly, we also observe that HiP-RSSM often exceeds traditional RSSMs in performance for dynamical systems previously assumed to have unchanging global dynamics due to the identification of unobserved causal effects in the dynamics.

§ RELATED WORK

Deep State Space Models.  Classical State-space models (SSMs) are popular due to their tractable inference and interpretable predictions. However, inference and learning in SSMs with high dimensional and large datasets are often not scalable. Recently, there have been several works on deep SSMs that offer tractable inference and scalability to high dimensional and large datasets. <cit.> use neural network architectures based on closed-form solutions of the forward inference algorithm on SSMs and achieve state-of-the-art results in state estimation and dynamics prediction tasks. <cit.> perform learning and inference in SSMs using variational approximations. However, most of these recurrent state-space models assume that the dynamics is fixed, which is a significant drawback, since this is rarely the case in real-world applications such as robotics.

Recurrent Switching Dynamical Systems. <cit.> tries to address the problem of changing/multimodal dynamics by incorporating additional discrete switching latent variables. However, these discrete states make learning and inference more involved.  <cit.> uses auxiliary variable methods for approximate inference in a multi-stage training process, while <cit.> uses variational approximations and relies on additional regularization/annealing to encourage discrete state transitions. On the other hand, <cit.> uses “soft” switches, which can be interpreted as a switching linear dynamical system which interpolate linear regimes continuously rather than using truly discrete states. We take a rather different, simpler formalism for modeling changing dynamics by viewing it as a multi-task learning problem with a continuous hierarchical hidden parameter that model the distribution over these tasks. Further our experiments in appendix <ref> show that our model significantly outperforms the soft switching baseline fraccaro2017disentangled.

Hidden Parameter MDPs.  Hidden Parameter Markov Decision Process (HiP-MDP) doshi2016hidden,killian2016transfer address the setting of learning to solve tasks with  similar, but not identical, dynamics. In HiP-MDP formalism, the states are assumed to be fully observed. However, we formalize the partially observable setting where we are interested in modelling the dynamics in a latent space under changing scenarios. The formalism is critical for learning from high dimensional observations and dealing with partial observability and missing data. The formalism can be easily extended to HiP-POMDPs by including rewards into the graphical model <ref>, for planning and control in the latent space hafner2019learning,sekar2020planning. However this is left as a future work.

Meta Learning For Changing Dynamics. There exists a family of approaches that attempt online model adaptation to changing dynamics scenarios via meta-learning nagabandi2018learning, nagabandi2018deep. They perform online adaptation on the

The graphical model for a particular instance for the HiP-SSM. The transition dynamics between latent states is a function of the previous latent state, and a specific latent task parameter $ \cvec{l} $ which is inferred from a context set of observed past observations. Actions are omitted for brevity.

parameters of the dynamics models through gradient descent finn2017model from interaction histories. Our method fundamentally differs from these approaches in the sense that we do not perform a gradient descent step at every time step during test time, which is computationally impractical in modern robotics, where we often deal with high-frequency data. We also empirically show that our approach adapts better, especially in scenarios with non-markovian dynamics, a property that is often encountered in real-world robots due to stiction, slip, friction, pneumatic/hydraulic/cable-driven actuators etc. <cit.> on the other hand, learn context-conditioned hierarchical dynamics model for control in a formalism similar to HiP-MDPs.  The former meta-learn a context conditioned gaussian process while the later use a context conditioned determisitic GRU. Our method on the other hand focuses on a principled probabilistic formulation and inference scheme for context conditioned recurrent state space models, which is critical for learning under partial observabilty/high dimensional observations, with noisy and missing data.

§ HIDDEN PARAMETER STATE SPACE MODELS (HIP-SSMS)

We denote a set of SSMs with transition dynamics $ f_{\cvec{l}} $  that are fully described by hidden parameters $ \cvec{l} $ and observation model $ h $  as a Hidden Parameter SSM (HiP-SSM). In this definition we assume the observation model $h$ to be independent of the hidden parameter $\cvec{l}$ as we only focus on cases where the dynamics change. HiP-SSMs allow us to extend SSMs to multitask settings where dynamics can vary across tasks. Defining the changes in dynamics by a latent variable unifies the dynamics across tasks as a single global function. In dynamical systems, for example, parameters can be physical quantities like gravity, friction of a surface, or the strength of a robot actuator. Their effects influence the dynamics but are not directly observed, so $\cvec{l}$ is not part of the observation space and is treated as a latent task parameter vector. The Bayesian network corresponding to HiP-SSM is shown in Figure <ref> and a formal definition is given below.

A HiP-SSM is represented by a tuple $\lbrace{\mathcal{L},\mathcal{C},\mathcal{Z},\mathcal{A},\mathcal{W},g,f,h\rbrace}$, where $\mathcal{Z}$, $\mathcal{A}$, $\mathcal{W}$, $\mathcal{L}$ and $\mathcal{C}$ denote the sets of hidden states $\cvec{z}$, actions $\cvec{a}$, observations $\cvec{w}$, latent tasks $\cvec{l}$ and contexts $C_l$ respectively.
Based on a task model $g$, transition model $f$ and observation model $g$, HiP-SSM have the following causal relationship:

\[
\begin{aligned}
\cvec{C}_l &= g(\cvec{l}) + \cvec{v}_l, \\
\cvec{z}_t &= f(\cvec{z}_{t-1},\cvec{a}_{t-1},\cvec{l}) + \cvec{\epsilon}_t, \\
\cvec{w}_t &= h(\cvec{z}_t) + \cvec{u}_t.
\end{aligned}
\]

Here, $\cvec{z}_t \in \mathcal{Z}$, $\cvec{a}_t \in \mathcal{A}$, and $\cvec{w}_t \in \mathcal{W}$ represent latent states, actions, and observations at time $t$, respectively. The vector $\cvec{\epsilon}_t$ denotes the process / transition noise, while $\cvec{u}_t$ and $\cvec{v}_l$ denote the measurement noise corresponding to observations and contexts.

Thus, a HiP-SSM describes a class of dynamics, and a particular instance of that class is obtained by fixing the parameter vector $\cvec{l} \in \mathcal{L}$. The dynamics $f$ for each instance depends on the value of the hidden parameters $\cvec{l}.$
Each instance of a HiP-SSM is an SSM conditioned on $\cvec{l}$. We also make the additional assumption that the parameter vector $\cvec{l}$ is fixed for the duration of the task (i.e., a local segment of a trajectory), and thus the latent task parameter has no dynamics. This assumption considerably simplifies the procedure to infer hidden parameterization and is reasonable, since dynamics can be assumed to be locally consistent over small trajectories in most applications nagabandi2018learning.

The definition is based on the related literature on HiP-MDPs doshi2016hidden, where the only unobserved variable is the latent task variable. One can connect HiP-SSMs with HiP-MDPs by including rewards in the definition and formalizing HiP-POMDPs. However, this is left for future work.

§ EXACT INFERENCE IN HIP-SSMS

The Gaussian graphical model corresponding to a particular instance $\cvec{l} \in \mathcal{L}$ for the HiP-RSSM. The posterior of the latent task/context variable is inferred via a Bayesian aggregation procedure described in <ref> based on a set of N interaction histories. The prior over the latent states $\cvec{z}_t^-$ is inferred via task conditional Kalman time update described in <ref> and the posterior over the latent states $\cvec{z}_t^+$ is inferred via Kalman measurement update described in <ref>. Here, the hollow arrows denote deterministic transformation leading to implicit distributions, using the context and observation encoders.

We perform inference and learning in the HiP-SSM borrowing concepts from both deep learning and graphical model communities following recent work on recurrent neural network models haarnoja2016backprop,becker2019recurrent, where the architecture of the network is informed by the structure of the probabilistic state estimator. We denote the resulting probabilistic recurrent neural network architecture as the Hidden Parameter Recurrent State Space Model (HiP-RSSM).

The structure of the Bayesian network shown in Figure <ref> allows a tractable inference of latent variables by the forward algorithm jordan2004graphical, koller2009probabilistic. Since we are dealing with continuous dynamical systems, we assume a Gaussian multivariate distribution over all variables (both observed and hidden) for the graph shown in Figure <ref>. This assumption has several advantages. Firstly, it makes the inference very similar to the well studied Kalman Filtering approaches. Secondly, the Gaussian assumptions and conditional in-dependencies allows us to have a closed form solution to each of these update procedures which are fully differentiable and can be backpropagated to the deep encoders. The update of beliefs about the hidden variables $\cvec{z}_t$ and $\cvec{l}$ takes place in three stages. Similar to Kalman filtering approaches, we have two recursive belief state update stages, the time update and observation update which calculate the prior and posterior belief over the latent states, respectively. However, we have an additional hierarchical latent variable $\cvec{l}$ which models the (unobserved) causal factors of variation in dynamics in order to achieve efficient generalization. Hence, we have a third belief update stage to calculate the posterior over the latent task variable based on the observed context set.  Each of these three stages is detailed in the following sections:

§.§ Inferring The Latent Task Variable (Context Update)

In this stage, we infer the posterior over the Gaussian latent task variable $\cvec{l}$ based on an observed context set $\mathcal{C}_{\cvec{l}}$. For any HiP-RSSM instance defined on a target sequence $\mathcal{T}=(\cvec{o}_t,\cvec{a}_t,\cvec{o}_{t+1},...,\cvec{o}_{t+N},\cvec{a}_{t+N})$, over which we intend to perform state estimation/prediction, we maintain a fixed context set $\mathcal{C}_{\cvec{l}}$. $\mathcal{C}_{\cvec{l}}$ in our HiP-SSM formalism can be obtained according to the algorithm designer's choice. We choose a $\mathcal{C}_{\cvec{l}}$ consisting of a set of tuples $\mathcal{C}_{\cvec{l}} = \lbrace \cvec{o}_{t-n}^{\cvec{l}},\cvec{a}_{t-n}^{\cvec{l}},\cvec{o}_{t-n+1}^{\cvec{l}}\rbrace_{n=1}^N$. Here each set element is a tuple consisting of the current state/observation, current action, and next state/observation for the previous N time steps.

Inferring a latent task variable $\cvec{l} \in \mathcal{L}$ based on an observed context set $\mathcal{C}_{\cvec{l}} = \{\cmat{x}_{n}^{\cvec{l}}\}_{n=1}^N$ has previously been explored by different neural process architectures gordon2018meta,garnelo2018neural.  Neural processes are multitask latent variable models that rely on deep set functions zaheer2017deep to generate a latent representation from a varying number of context points in a permutation invariant manner.
Similar to <cit.> we formulate the aggregation of context data as a Bayesian inference problem, where the information contained in $\mathcal{\cmat{C}}_l$ is directly aggregated into the statistical description of $\cvec{l}$ based on a factorized Gaussian observation model of the form $p(\cvec{r}_n^{\cvec{l}}|\cvec{l})$, where

\begin{align*}
\cvec{r}_n^{\cvec{l}} & = \textrm{enc}_{\cvec{r}}(\cvec{o}_{t-n}^{\cvec{l}},\cvec{a}_{t-n}^{\cvec{l}},\cvec{o}_{t-n+1}^{\cvec{l}}), \\ \cvec{\sigma}_n^{\cvec{l}} & = \textrm{enc}_{\cvec{\sigma}}(\cvec{o}_{t-n}^{\cvec{l}},\cvec{a}_{t-n}^{\cvec{l}},\cvec{o}_{t-n+1}^{\cvec{l}}).
\end{align*}

Here $n$ is the index of an element from the context set $\mathcal{\cvec{C}}_{\cvec{l}}$.

Given a prior $p_0(\cvec{l}) = \mathcal{N}(\cvec{l}|\cvec{\mu}_{0},\textrm{diag}(\cvec{\sigma}_{0}))$ we can compute the posterior $p(\cvec{l}|\mathcal{\cmat{C}}_{\cvec{l}})$ using the Bayes rule.

The Gaussian assumption allows us to obtain a closed-form solution for the posterior estimate of the latent task variable, $p(\cvec{l}|\mathcal{\cmat{C}}_{\cvec{l}})$ based on Gaussian conditioning. The factorization assumption further simplifies this update rule by avoiding computationally expensive matrix inversions into a simpler update rule as

\begin{align*}
\cvec{\sigma_{l}} & = \left( \left(\cvec{\sigma}_{0}\right)^\ominus+ \sum_{n=1}^N \left(\cvec{\sigma}_n^{\cvec{l}}\right)^\ominus\right)^\ominus, \\ \cvec{\mu}_{\cvec{l}} & = \cvec{\mu}_{0} + \cvec{\sigma}_{\cvec{l}} \odot \sum_{n=1}^N \left(\cvec{r}^{\cvec{l}}_n - \cvec{\mu}_{0}\right) \oslash \cvec{\sigma}_n^{\cvec{l}}
\end{align*}

where $\ominus$, $\odot$ and $\oslash$ denote element-wise inversion, product, and division, respectively. Intuitively, the mean of the latent task variable $\cvec{\mu}_{\cvec{l}}$ is a weighted sum of the individual latent observations $\cvec{r}^{\cvec{l}}_n$, while the variance of the latent task variable $\cvec{\sigma}_{\cvec{l}}$ gives the uncertainty of this task representation.

§.§ Inferring Prior Latent States over the Next Time Step (Task Conditional Prediction)

The goal of this step is to compute the prior marginal

\begin{align}
\label{eq:predict}
\begin{split}
p(\cvec{z}_t|\cvec{w}_{1:t-1},\cvec{a}_{1:t},\mathcal{C}_{\cvec{l}}) = \iint p(\cvec{z}_t|\cvec{z}_{t-1},\cvec{a}_{t},\cvec{l})&p(\cvec{z}_{t-1}|\cvec{w}_{1:t-1},\cvec{a}_{1:t-1},\mathcal{C}_{\cvec{l}}) \\ &p(\cvec{l}|\mathcal{C}_{\cvec{l}})d\cvec{z}_{t-1}d\cvec{l}.
\end{split}
\end{align}

We assume a dynamical model of the following form: $p(\cvec{z}_t|\cvec{z}_{t-1},\cvec{a}_{t},\cvec{l}) = \mathcal{N}\left(\cmat{A}_{t-1}\cvec{z}_{t-1} + \cmat{B}\cvec{a}_t + \cmat{C}\cvec{l},  \cmat{\Sigma}_{\textrm{trans}}\right)$ and
denote the posterior from the previous time-step by $p(\cvec{z}_{t-1}|\cvec{w}_{1:t-1},\cvec{a}_{1:t-1},\mathcal{C}_{\cvec{l}}) = \mathcal{N}\left(\cvec{z}_{t-1}^+, \cmat{\Sigma}_{t-1}^+\right)$.

Following <cit.>, we assume that the action $\cvec{a}_t$ is known and not subject to noise.

At any time t, the posterior over the belief state $\cvec{z}_{t-1}|\cvec{w}_{1:t-1},\cvec{a}_{1:t-1}$, posterior over the latent task variable $\cvec{l}|\mathcal{C}_{\cvec{l}}$ and the action $\cvec{a}_t$ are independent of each other since they form a “common effect" / v structure koller2007graphical with the unobserved variable $\cvec{z}_{t}$.  With these independencies and Gaussian assumptions, according to the Gaussian identity <ref>, it can be shown that calculating the integral in equation <ref> has a closed form solution as follows,
$p(\cvec{z}_t|\cvec{w}_{1:t-1},\cvec{a}_{1:t},\mathcal{C}_{\cvec{l}}) = \mathcal{N}(\cvec{z}_t^-,\cmat{\Sigma}_t^-)$, where

\begin{align}
\label{eq:cond-predict}
\cvec{z}_t^- &=\cmat{A}_{t-1}\cvec{z}_{t-1} + \cmat{B}\cvec{a}_t + \cmat{C}\cvec{l}, \\
\cmat{\Sigma}_t^- &= \cmat{A}_{t-1}\cmat{\Sigma}_{t-1}^+\cmat{A}_{t-1}^T + \cmat{C}\cmat{\Sigma}_{\cvec{l}}\cmat{C}^T + \cmat{\Sigma}_{\textrm{trans}}.
\end{align}

[Linear Combination Gaussian Marginalization]gigausslinearmarg
If $\cvec{u} \sim \mathcal{N}(\cvec{\mu}_{\cvec{u}} + \cvec{b},\cmat{\Sigma}_{\cvec{u}})$ and $\cvec{v} \sim \mathcal{N}(\cvec{\mu}_{\cvec{v}},\cmat{\Sigma}_{\cvec{v}})$  are normally distributed independent random variables and if the conditional distribution $p\left(\cvec{y}|\cvec{u}, \cvec{v}\right) = \mathcal{N}(\cmat{A}\cvec{u} + \cvec{b} + \cmat{B}\cvec{v}, \cmat{\Sigma})$, then marginal $p(\cvec{y}) = \int p(\cvec{y}|\cvec{u},\cvec{v})p(\cvec{u})p(\cvec{v})) d\cvec{u}d\cvec{v} = \mathcal{N}(\cmat{A}\cvec{\mu}_{\cvec{u}} + \cvec{b} + \cmat{B}\cvec{\mu}_{\cvec{v}}, \cmat{A}\cmat{\Sigma}_u\cmat{A}^T + \cmat{B}\cmat{\Sigma}_{\cvec{v}}\cmat{B}^T +  \cmat{\Sigma}).$

The proof for the above identity is given in Appendix.

Modelling Choices We use a locally linear transition model $\cmat{A}_{t-1}$ as in <cit.> (see Appendix <ref>) and a nonlinear control model as in <cit.> (Chapter <ref>). The local linearization around the posterior mean, can be interpreted as equivalent to an EKF. For the latent task transformation we can either use a linear, locally-linear or non-linear transformation. More details on the latent task transformation model can be found in the Appendix <ref>. Our estimations (Figure <ref>) show that a nonlinear feedforward neural network $f(.)$ that outputs mean and variances and interacts additively gave the best performance in practice.  $f(.)$ transforms the latent task moments $\cvec{\mu_l}$ and $\cvec{\sigma_l}$ directly into the latent space of the state-space model through additive interactions. The corresponding time update equations are given below:

\begin{align*}
\cvec{z}_t^- &=\cmat{A}_{t-1}\cvec{z}_{t-1}^+ + \cmat{b}(\cvec{a}_t) + \cmat{f}(\cvec{\mu}_l), \\
\cmat{\Sigma}_t^- &= \cmat{A}_{t-1}\cmat{\Sigma}_{t-1}^+\cmat{A}_{t-1}^T + \cmat{f}(\cvec{\sigma}_{\cvec{l}})+ \cmat{\Sigma}_{\textrm{trans}}.
\end{align*}

This chapter primarily aimed at accurately modeling mean predictions under non-stationary conditions, focusing on minimizing mean squared errors. We opted for locally linear transitions $\cvec A_t$ and non-linear transformations $f$ to enhance mean predictions, without prioritizing uncertainty quantification. In contrast, Chapter <ref> equally emphasizes uncertainty quantification and mean predictions. Subsequent chapters adopted a simple linear model for both matrices $\cvec A$ and $\cvec C$, chosen for its robustness in both mean and uncertainty predictions and computational efficiency gu2021efficiently,mondal2023efficient.

§.§ Inferring posterior latent states / Observation Update

The goal of this step is to compute the posterior belief $p(\cvec{z}_t|\cvec{o}_{1:t},\cvec{a}_{1:t},\cvec{C_l})$. We first map the observations at each time step $\cvec{o}_t$ to a latent space using an observation encoder haarnoja2016backprop,becker2019recurrent that emits latent features $\cvec{w}_t$ along with uncertainty in those features via a variance vector $\cvec{\sigma}_{obs}^t$. We then computed the posterior belief $p(\cvec{z}_t|\cvec{w}_{1:t},\cvec{a}_{1:t},\cvec{C_l})$, based on $p(\cvec{z}_t|\cvec{w}_{1:t-1},\cvec{a}_{1:t},\cvec{C_l}) $ obtained from the time update, the latent observation $(\cvec{w}_t,\cvec{\sigma}_t^{obs})$ and the observation model $\cmat{H}$. This is exactly the Kalman update step, which has a closed form solution as shown below for a time instant $t$,

\begin{align*}
&\text{Kalman Gain:}  &\cmat{Q}_t &= \cmat{\Sigma_t}^- \cmat{H}^T\left(\cmat{H} \cmat{\Sigma_t}^- \cmat{H}^T  + \cvec{I} \cdot \cmat{\sigma_t}^\mathrm{obs} \right)^{-1}, \\
&\text{Posterior Mean:}  &\cvec{z_t}^+ &= \cvec{z_t}^- + \cmat{Q}_t \left(\cvec{w_t} - \cmat{H} \cvec{z_t}^-  \right), \\
&\text{Posterior Covariance:}  &\cmat{\Sigma_t}^+ &= \left(\cmat{I} - \cmat{Q}_t \cmat{H} \right) \cmat{\Sigma_t}^-,
\end{align*}

where $\cmat{I}$ denotes the identity matrix. This update is added as a layer in the computation graph (<cit.>; <cit.>).  However, the Kalman update involves computationally expensive matrix inversions of the order of $\mathcal{O}(L^3)$, where $L$ is the dimension of the latent state $\cvec{z}_t$. Thus, in order to make the approach scalable, we follow the same factorization assumptions as in <cit.>. This factorization provides a simple way to reduce the observation update equation to a set of scalar operations, reducing the computational complexity from $\mathcal{O}(L^3)$ to $\mathcal{O}(L)$.
More mathematical details on the simplified update equation can be found in Section <ref>.
From a computational perspective, this a Gaussian conditioning layer, similar to section <ref>. Both output a posterior distribution on latent variables $\cvec{z}$, given a prior $p(\cvec{z})$ and an observation model $p(\boldsymbol{w}|\boldsymbol{z})$, using the Bayes rule: $p(\boldsymbol{z}|\boldsymbol{w}) = p(\boldsymbol{w}|\boldsymbol{z})p(\boldsymbol{z})/p(\boldsymbol{w})$. This has a closed form solution because of Gaussian assumptions, which is coded as a layer in the neural network. The observation model is assumed to have the following structure, $P(\boldsymbol{w}|\boldsymbol{z}) = \mathcal{N}(\boldsymbol{H}\boldsymbol{z}, \boldsymbol{\Sigma}_{obs})$.

§ HIP-SSM AS ADAPTIVE MULTI-TASK WORLD MODELS

Using adaptive internal world models to envision the future: Our predictions about future world states rely on more than just a sequence of action signals; they are also conditioned on latent task variables ($l$). These latent variables introduce additional causal factors that influence the dynamics, factors not considered in basic State Space Model (SSM) approaches introduced in Chapter <ref>. This inclusion enables a deeper and more nuanced ability to manipulate and foresee future states in diverse scenarios.

§.§ End To End Learning For Online Adaptation To Changes

Multitask Data CreationThe latent task variable $l$ models a distribution over functions garnelo2018neural, rather than a single function. In our case, these functions are latent dynamics functions. In order to train such a model, we use a training procedure that reflects this objective, where we form datasets consisting of timeseries, each with different latent transition dynamics. Thus, we collect a set of trajectories over which the dynamics changes over time. These can be trajectories where a robot picks up objects of different weights or a robot traverses terrain of different slopes. Now, we introduce a multitask setting with a rather flexible definition of task, where each temporal segment of a trajectory can be considered to be a different “task" and the observed context set based on interaction histories from the past $N$ time steps provides information about the current task setting. This definition allows us to have a potentially infinite number of tasks/local dynamical systems, and the distribution over these tasks/systems is modeled using a hierarchical latent task variable $l$. The formalism is based on the assumption that over these local temporal segments the dynamics is unchanging. This local consistency in dynamics holds for most real world scenarios nagabandi2018deep,nagabandi2018learning. However, our formalism can model the global changes in dynamics at the test time, since we obtain a different instance of the HiP-RSSM for each temporal segment based on the observed context set. We also provide a detailed description of the multitask data set creation process in Table <ref> and a pictorial illustration in Appendix <ref>.

Multi Task Dataset Creation For Training HiP-RSSM
Required:A set $S$ of trajectories of changing dynamics

$D \leftarrow \phi$

each trajectory $\tau \in S$

* divide the trajectory $\tau$ into non-overlapping windows $T_l$
of length N. Let $T = \{T_1, T_2, T_3, .,.,.\}$ be the list of all temporal segments/time-series.

each time window $T_l \in T$

* maintain a context set $C_l$ consisting of N previous interactions;

* update $D \leftarrow D \cup \{C_l,T_l\}$

Output: Output $D$ consisting of batch of context and target sets.

Batch Training. Let $T  \in \mathcal{R}^{B \times N \times D}$ be the batch of local temporal segments with different dynamics that we intend to model with the HiP-RSSM formalism. Given a target batch $T$, HiP-RSSM can be trained in a supervised manner similar to popular recurrent architectures such as LSTMs or GRUs.  However, for each local temporal sequence $t \in T$, in which the dynamics is modeled, we also input a set of previous interactions $N$, which forms the context set $C \in \mathcal{R}^{B \times N \times D}$ to infer the latent task as explained in Section <ref>. Processing the context set $C$ adds minimal additional computational / memory constraints, as we use a permutation-invariant set encoder. The set encoder allows for parallelization in processing the context set as opposed to recurrent processing of the target set.

The learnable parameters in the computation graph include the locally linear transition model $\cvec{A}_t$, the nonlinear control factor $\cvec{b}$, the linear/nonlinear latent task transformation model $C$, the transition noise $\cvec{\Sigma}_{\textrm{trans}}$, along with the observation encoder, context encoder, and output decoder.

Loss Function. The network is tasked to minimize the prediction errors by maximizing the posterior predictive log-likelihood, which is given below for a single trajectory, i.e.,

\begin{align}
\label{eq:objective}
L & =\sum_{t=1}^H \log p(\cvec{o}_{t+1}|\cvec{w}_{1:t}, \cvec{a}_{1:t}, \cvec C_l) \nonumber \\ & = \sum_{t=1}^H \log \int p(\cvec{o}_{t+1}|\cvec{z}_{t+1})  p(\cvec{z}_{t+1}|\cvec{w}_{1:t}, \cvec a_{1:t}, \cvec C_l) d\cvec{z}_{t+1} \nonumber \\
\end{align}

The extension to multiple trajectories is straightforward and was omitted to keep the notation uncluttered. Here, $\cvec{o}_{t+1}$ are the ground truth observations at the time step $t+1$ that must be predicted from all observations up to the time step $t$.

Approximating the likelihood We employ a Gaussian approximation of the posterior predictive log-likelihood of the form $ p(\cvec{o}_{t+1}|\cvec{w}_{1:t}, \cvec{a}_{1:t}, \cvec C_l) \approx \mathcal{N}(\cvec{\mu}_{\cvec{o}_{t+1}},\textrm{diag}(\cvec{\sigma}_{\cvec{o}_{t+1}}))$ where we use the mean of the prior belief $\cvec{\mu}_{z_{t+1}}^-$ to decode the predictive mean, that is, $\cvec{\mu}_{\cvec{o}_{t+1}} =  \textrm{dec}_{\cvec{\mu}}(\cvec{\mu}_{z_{t+1}}^{-})$ and the variance estimate of the prior belief to decode the observation variance, that is, $\cvec{\sigma}_{o_{t+1}} = \textrm{dec}_{\sigma}(\cmat{\Sigma}_{z_{t+1}}^{-})$. This approximation can be motivated by a moment matching perspective and allows end-to-end optimization of logarithmic likelihood without using auxiliary objectives such as the ELBO <cit.>. Thus, the approximate Gaussian predictive log-likelihood for a single sequence is then computed as

\begin{align}
\label{eq: likeli-acrkn}
&\mathcal{L}\left(\cvec{o}_{(1:T)}\right) = \\&
\dfrac{1}{T} \sum_{t=1}^T  \log \mathcal{N}\left(\cvec{o}_t \bigg| \textrm{dec}_{\mu}(\cvec{z}_t^+), \textrm{dec}_{\Sigma} (\cvec{\sigma}^{\mathrm{u},+}_t,
\cvec{\sigma}^{\mathrm{s},+}_t,
\cvec{\sigma}^{\mathrm{l},+}_t )\right), \nonumber
\end{align}

where $\textrm{dec}_{\mu}(\cdot)$ and $\textrm{dec}_{\Sigma}(\cdot)$ denote the parts of the decoder that are responsible for decoding the latent mean and latent variance respectively.

Variations Of The Learning Objectives

We optimize the root mean square error (RMSE) between the decoder output and the ground-truth states. As in <cit.> we use the differences to the next state as our ground truth states, as this results in better performance for dynamic learning, especially at higher frequencies. In principle, we could train on the Gaussian log-likelihood instead of the RMSE and hence model the variances. Training in RMSE yields slightly better predictions and allows a fair comparison with deterministic baselines that use the same loss, such as feedforward neural networks, LSTMs and metalearning algorithms such as MAML finn2017model. Therefore, we report results with the RMSE loss.

Depicts the schematic of HiP-RSSM. The output of the task conditional `Time Update' stage, which forms the prior for the next time step $(z_{t+1}^-,\Sigma_{t+1}^-)$ is decoded to get the prediction of the next observation.

Gradients are computed using (truncated) backpropagation through time (BPTT) (<cit.>) and clipped. We optimize the objective using the Adam (<cit.>) stochastic gradient descent optimizer with default parameters.

Irrespective of the learning objectives/loss functions used, we still have to model uncertainty in HiP-SSM latent state as the context set and observations do not contain the full task and state information of the system.
The architecture of the HiP-SSM is summarized in Figure <ref>.

§.§ HiP-RSSM during Test Time / Inference

We perform inference using HiP-RSSM at test time on a trajectory with varying dynamics using algorithm <ref>. A pictorial representation of the same is given in the Figure <ref>. We use this inference scheme to visualize how the latent variable $l$, that describe different instances of a HiP-RSSM over different temporal segments, evolve at a global level. The visualizations are reported in Figures <ref> and <ref> in the experiments section discussed below.

HiP-RSSM Test Time Inference

Required: Trained HiP-RSSM Model

Required: A time series $\tau$ of length $K>>N$

Divide the time series $\tau$ into non-overlapping windows $T_l$ of length N. Let $T = \{T_1, T_2, T_3, .,.,.\}$ be the ordered list of all temporal segments, sorted in the ascending order of time of occurrence.

each time window $T_l \in T$

* maintain a context set $C_l$ consisting of N previous interactions;

* infer posterior latent task variable $p(\cvec{l}|\mathcal{C}_{\cvec{l}})$ using context update stage as in section <ref>;

* using the posterior over latent task variable $\cvec{l}|\mathcal{C}_{\cvec{l}}$ and observations in sequence $T_l$ to perform sequential Bayesian inference over the state space model using Kalman observation update (<ref>) and task conditional Kalman time update; (<ref>)

HiP-RSSM Inference Under Changing Dynamics Scenarios

§ EXPERIMENTS

This section evaluates our approach on a diverse set of dynamical systems from the robotics domain in simulations and real systems. We show that HiP-RSSM outperforms contemporary recurrent state-space models (RSSMs) and recurrent neural networks (RNNs) by a significant margin under changing dynamics scenarios. Further, we show that HiP-RSSM outperforms these models even under situations with partial observability/ missing values. We also baseline our HiP-RSSM with contemporary multi-task models and improve performance, particularly in modelling non-Markovian dynamics and under partial observability. Finally, the visualizations of the Gaussian latent task variables in HiP-RSSM demonstrates that they learn meaningful representations of the causal factors of variations in dynamics in an unsupervised fashion.

We consider the following baselines:

* RNNs - We compare our method to two widely used recurrent neural network architectures, LSTMs lstm and GRUs gru.

* RSSMs - Among several RSSMs from the literature, we chose RKN <cit.> as these have shown excellent performance for dynamics learning shaj2020action and relies on exact inference as in our case.

* Multi Task Models - We also compare with state of the art multi-task learning models like Neural Processes (NPs) and MAML nagabandi2018learning. Both models receive the same context information as in HiP-RSSM.

In case of recurrent models we replace the HiP-RSSM cell with a properly tuned LSTM, GRU and RKN Cell respectively, while fixing the encoder and decoder architectures. For the NP baseline we use the same context encoder and aggregation mechanism as in HiP-RSSM to ensure a fair comparison. We create partially observable settings by imputing $50\%$ of the observations during inference. More details regarding the baselines and hyperparameters can be found in the Appendix <ref>.

§.§ Soft Robot Playing Table Tennis

We first evaluate our model on learning the dynamics of a pneumatically actuated muscular robot. This four Degree of Freedom (DoF) robotic arm is actuated by Pneumatic Artificial Muscles (PAMs) buchler2016lightweight. The data consists of trajectories of hitting movements with varying speeds while playing table tennis buchler2020learning. This robot's fast motions with high accelerations are complicated to model due to hysteresis and hence require recurrent models shaj2020action.

We show the prediction accuracy in the RMSE in Table <ref>. We observe that the HiP-RSSM can outperform the previous state of the art predictions obtained by recurrent models. Based on domain knowledge, we hypothesize that the latent context variable captures multiple unobserved causal factors of variation that affect the dynamics in the latent space, which are not modeled in contemporary recurrent models. These causal factors could be, in particular, the temperature changes or the friction due to a different path that the Bowden trains take within the robot. Disentangling and interpreting these causal factors can be exciting and improve generalization, but it is out of scope for the current work. Further, we find that the multitask models like NPs and MAML fail to model these dynamics accurately compared to all the recurrent baselines because of the non-markovian dynamics resulting from the high accelerations in this pneumatically actuated robot.

No Imputation      $50 \%$ Imputation

HiP-RSSM     $2.30 \pm0.043$     $2.47\pm 0.012$

RKN        $3.088\pm 0.046$       $3.223\pm0.014$

LSTM     $3.108\pm 0.041$      $3.630\pm0.097$

GRU     $3.287\pm 0.013$      $3.621\pm0.047$

FFNN    $8.150\pm 0.047$       -

NP     $5.526\pm 0.030$      -

MAML    $7.314 \pm 0.021$        -

Pneumatic RMSE $\left(10^{-3}\right)$

No Imputation      50 $\%$ Imputation
HiP-RSSM    $2.833 \pm 0.024$     $2.843 \pm 0.024$

RKN         $3.392 \pm 0.062$        $ 3.398 \pm 0.062$

LSTM     $3.503 \pm0.006$      $ 3.736 \pm 0.062$

GRU     $3.407 \pm 0.02$     $ 3.642 \pm 0.153$

FFNN     $3.313 \pm 0.018 $      -

NP     $2.765 \pm0.004 $      -

MAML     $3.202 \pm 0.006$     -

Franka RMSE $\left(10^{-4}\right)$

Prediction Error in RMSE for (a) pneumatic muscular arm (<ref>) and (b) Franka Arm manipulating varying loads (<ref>) for both fully observable and partially observable scenarios.

§.§ Robot Manipulation With Varying Loads

We collected data from a $7$ DoF Franka Emika Panda manipulator carrying six different loads at its end effector. It involved a mix of movements of different velocities from slow to swift movements that covered the entire robot workspace. We chose trajectories with four different loads as training set and evaluated the performance on two unseen weights, which results in a scenario where the dynamics change over time. Here, the causal factor of variation in dynamics is the weights attached to the end-effector and assumed to be unobserved.
We show the prediction errors in RMSE in Table <ref>. HiP-RSSMs outperforms all recurrent state-space models, including RKN and deterministic RNNs, in modelling these dynamics under fully observable and partially observable conditions. The multi-task baselines of NPs and MAML perform equally well under full observability for this task because of the near Markovian dynamics of Franka Manipulator, which often does not need recurrent models. However, HiP-RSSMs have an additional advantage in that these are naturally suited for partially observable scenarios and can predict ahead in a compact latent space, a critical component for recent success in model-based control and planning hafner2019learning.

(a) Simulated environment of a Wheeled Robot traversing terrains of different slopes. (b) and (c) shows the tSNE van2008visualizing plots of the latent task embeddings produced from randomly sampled instances of HiP-RSSM for two different robots. (b) The wheeled robot discussed in section <ref> traversing terrains of varying slopes. The color coding indicates the average gradients of the terrain for each of these instances. These can have either positive or negative values. (c) The Franka manipulator with loads of varying weights attached to the end-effector. The color coding indicated weights ranging from 0 to 3 kilograms.

§.§ Robot Locomotion In Terrains Of Different Slopes

Wheeled mobile robots are the most common types of robots being used in exploration of unknown terrains where they may face uneven and challenging terrain. We set up an environment using a Pybullet simulator coumans2016pybullet where a four-wheeled mobile robot traverses an uneven terrain of varying steepness generated by sinusoidal functions sonker2020adding as shown in <ref>. This problem is challenging due to the highly non-linear dynamics involving wheel-terrain interactions. In addition, the varying steepness levels of the terrain results in a changing dynamics scenario, which further increases the complexity of the task.

We show the prediction errors in RMSE in Table <ref>. When most recurrent models, including RSSMs and deterministic RNNs, fail to model these dynamics, HiP-RSSMs are by far the most accurate in modelling these challenging dynamics in the latent space. Further, the HiP-RSSMs perform much better than state of the art multi-task models like NPs and MAML.

No Imputation     $50\%$ Imputation

HiP-RSSM     $2.96 \pm 0.212$     $6.15 \pm 0.327$

RKN     $7.17\pm0.017$     $14.66 \pm 0.224$

LSTM     $9.14\pm0.026$     $51.21 \pm 0.431$

GRU     $9.216\pm0.073$     $53.14 \pm 0.242$

FFNN     $8.72\pm0.021$     -

NP     $4.57\pm0.013$     -

MAML     $5.04\pm0.051$     -

Prediction error for wheeled mobile robot trajectories in RMSE ($10^{-5}$) for both fully observable and partially observable scenarios.