Learning World Models With Hierarchical Temporal Abstractions: A Probabilistic Perspective
Vaisakh Shaj


Partial Renormalization of Quasiparticle Interactions
    Kun Chen
    April 29, 2024
=====================================================











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. 

 

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.


 
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 Y given Z, denoted X⊥Y|Z), for some subsets of variables X, Y, 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.




 §.§ 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.




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.




 





§ STATE SPACE MODELS


[10]r0.5
  


 !


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 ( 𝒵, 𝒜, 𝒪, f, h, Δ t), where 𝒵 is the state space, 𝒜 the action space and 𝒪 the observation space of the SSM. The parameter Δ 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

    [c]
      z_t    = f(z_t-1,a_t-1) + ϵ_t,   ϵ_t ∼𝒩(0, Σ^trans), 
    o_t    = h(z_t) + v_t,   v_t ∼𝒩(0, Σ^obs).

Here z_t ∈𝒵, a_t ∈𝒜 and o_t ∈𝒪 are the latent states, actions, and observations at time t. The vectors ϵ_t and 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 A is used to infer the current prior state estimate p(z_t|w_1:t-1) = 𝒩(z_t^-,Σ_t^-), i.e., a-priori to the observation, from the previous posterior estimate p(z_t-1|w_1:t-1) = 𝒩(z_t-1^+, Σ_t-1^+ ),  by 

    z^-_t   = Az^+_t-1
    and  Σ^-_t   = AΣ_t-1^+ A^T + Σ^trans.



  
Update Step
The prior estimate is then updated using the current observation w_t and the observation model H to obtain the posterior estimate (z_t^+, Σ_t^+ ), i.e.,

    z_t^+    = z_t^- + Q_t (w_t - Hz_t^-  ) 
    and  Σ^+_t    = (I - Q_t H) Σ^-_t 
    , with  Q_t    = Σ^-_t H^T(HΣ^-_t H^T  + Σ^obs)^-1,

where I denotes the identity matrix.
The matrix 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., 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 o to a latent observation w_t = enc_w( o_t). In addition to the latent observation, the encoder also learns to output the uncertainty of this observation, i.e., σ_o_t = enc_σ( 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) 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:


    
  * 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( z_t| o_1:t,  a_1:t-1) = 𝒩( z_t^+, Σ_t^+) such that only the diagonal and one off-diagonal vector of the covariance need to be computed, i.e. 

    Σ_t^+ = [[ Σ_t^u,+ Σ^s,+_t; Σ^s,+_t Σ^l,+_t ]],  with Σ_u^+ = diag(σ_t^u,+),  Σ_l^+ = diag(σ_t^l,+)

and Σ_s^+ = diag(σ_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:

    z_t^+ = z_t^- +
    [[ σ^u,-_t; σ^l,-_t ]]
    ⊙[[ w_t - z^u,-_t; w_t - z^u,-_t ]]
    ⊘[[ σ_t^u,- + σ_t^obs; σ_t^u,- + σ_t^obs ]],


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

    σ^u,+_t    = σ^u,-_t ⊙σ^u,-_t ⊘( σ_t^u,- + σ_t^obs),  
    σ^s,+_t    = σ^u,-_t ⊙σ^s,-_t ⊘( σ_t^u,- + σ_t^obs), 
    σ^l,+_t    = σ^l, -_t - σ^s,-_t ⊙σ^s,-_t ⊘( σ_t^u,- + σ_t^obs),

where ⊙ denotes the elementwise vector product and  ⊘ 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 H = [ I,   0] allows for the division of the latent state vector into two distinct components. The latent state z_t = [ p_t^T,  d_t^T]^T has twice the dimensionality of the latent observation w_t and only the first half of the latent state, i.e., p_t, can be observed. The upper part utilizes the identity matrix in H = [ I,   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., 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
[12]r.35
[b]0.31
   width=
         
     
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( l|  r_1:N) using the Bayes rule, given an observation model of the form, 
    p(r_n|l) = 𝒩(r_n| H l,diag(σ_n))
 with H =  I and a prior p( l) = 𝒩(μ_0, diag(σ_0)). The Gaussian assumption allows us to obtain a closed-form solution for the posterior estimate of the latent task variable p(l|𝒞_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,

    σ_l = ( (σ_0)^⊖+ ∑_n=1^N (σ_n)^⊖)^⊖,   μ_l = μ_0 + σ_l⊙∑_n=1^N (r_n - μ_0) ⊘σ_n.


Here ⊖, ⊙ and ⊘ 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 r_n and its corresponding uncertainty σ_o_n. We learn these using a Deep Set encoder as follows,

    p(r_n|l) = 𝒩(r_n|l,diag(σ_n)),   r_n = enc_r(x_n),   σ_n = enc_σ(x_n).



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





  
Bayesian Aggregation As Probabilistic Attention Intuitively the mean of the latent task variable μ_l is a uncertainty weighted sum of the individual latent observations r_n, while the variance of the latent task variable σ^2_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, H 
 =  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
[26]r0.35
  
    < g r a p h i c s >

  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(h_i) and the likelihood that the evidence e would occur, given h_i is P(e | h_i), the brain computes the posterior probability of hypothesis h_i, given the evidence e: P(h_i | e) using internal generative models. The figure shows simplified version of Bayes' rule that puts it together: P(h_i | e)=P(e | h_i) P(h_i).
  

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.





  
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
  
    < g r a p h i c s >

  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.


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
  
    < g r a p h i c s >

  (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







To achieve action conditioning within the recurrent cell, we modify the transition model proposed in <cit.> to include  a control model b⃗(a⃗_t) in addition to the locally linear transition model A⃗_t. 
The control model b⃗(a⃗_t) is added to the locally
linear state transition, i.e., z⃗^-_t+1 = A⃗_t z⃗^+_t + b⃗(a⃗_t). 
As illustrated in <ref>, this formulation of latent dynamics captures the causal effect of the actions variables 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 b⃗(a⃗_t) can be represented in several ways, i.e.:

  (i) Linear: b⃗_l(a⃗_t) =B⃗a⃗_t, where B⃗ is a linear transformation matrix.

  (ii) Locally-Linear: b⃗_m(a⃗_t) = B⃗_t a⃗_t, where B⃗_t = ∑_k=0^K β^(k)(z⃗_⃗t⃗) B⃗^(k) is a linear combination of k linear control models B⃗^(k). A small neural network with softmax output is used to learn β^(k).  

  (iii) Non-Linear: b⃗_n(a⃗_t) = f⃗(a⃗_t), where 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., Σ⃗^-_t+1 = A⃗_t Σ⃗_t^+ A⃗_t^T + I⃗·σ⃗^⃗t⃗r⃗a⃗n⃗s⃗ as the action is known and not uncertain.
Thus, unlike the state transition model A⃗_t, we do not need to constrain the model 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, b⃗_n(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



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: o⃗_1:t, a⃗_1:t↦o⃗_t+1 that predicts the next observation o⃗_t+1 given the histories of observation o⃗_1:t and actions 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 (Σ_t) about the belief of our future world states in addition to the mean estimates (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(z⃗_t+1| w⃗_1:t, 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., 

    L     =∑_t=1^H log p(o_t+1|w_1:t, a_1:t) 
        = ∑_t=1^H log∫ p(o_t+1|z_t+1)  p(z_t+1|w_1:t,  a_1:t) dz_t+1

The extension to multiple trajectories is straightforward and omitted to keep the notation uncluttered. Here, 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(o_t+1|w_1:t, a_1:t) ≈𝒩(μ_o_t+1,diag(σ_o_t+1)) where we use the mean of the prior belief μ_z_t+1^- to decode the predictive mean, i.e, μ_o_t+1 =  dec_μ(μ_z_t+1^-) and the variance estimate of the prior belief to decode the observation variance, i.e., σ_o_t+1 = dec_σ(Σ_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  

    ℒ(o_(1:T)) = 
       1T∑_t=1^T  log𝒩(o_t | dec_μ(z_t^+), dec_Σ (σ^u,+_t,
    σ^s,+_t,
    σ^l,+_t )),

where dec_μ(·) and dec_Σ(·) denote the parts of the decoder that are responsible for decoding the latent mean and latent variance respectively.



  
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.
    
    ℒ_fwd = √(1T∑_t=1^T  ∥(o⃗_t+1 -  o⃗_t) - dec(z⃗^-_t+1) ∥^2).
 
    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., 1ms, 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 ô⃗̂_t+1 = o⃗_t + dec(z⃗_t+1^-), where o⃗_t is the true observation in t and dec(z⃗_t+1^-) is the output of the actual decoder network. During inference, we introduce an additional memory m⃗_t that stores the last prediction and uses it as an observation in case of a missing observation, that is, ô⃗_t+1 = m⃗_t + dec(z⃗_t+1^-) where we use the true observation, m⃗_t= o⃗_t if available, and the predicted observation, m⃗_t= ô⃗̂_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 z⃗^-_t+1 to predict o⃗_t+1. 
The prior mean z⃗^-_t+1 integrates all information up to time step t, including 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 
 

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. 


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 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.


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.
  


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.
[11]r0.5 
  
    < g r a p h i c s >

  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>.





§ 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. 

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
[21]r.35
 0.75
[b].42
    
 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 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_l  that are fully described by hidden parameters 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 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 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 {ℒ,𝒞,𝒵,𝒜,𝒲,g,f,h}, where 𝒵, 𝒜, 𝒲, ℒ and 𝒞 denote the sets of hidden states z, actions a, observations w, latent tasks 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:

    C_l    = g(l) + v_l, 
    z_t    = f(z_t-1,a_t-1,l) + ϵ_t, 
    w_t    = h(z_t) + u_t.

Here, z_t ∈𝒵, a_t ∈𝒜, and w_t ∈𝒲 represent latent states, actions, and observations at time t, respectively. The vector ϵ_t denotes the process / transition noise, while u_t and 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 l∈ℒ. The dynamics f for each instance depends on the value of the hidden parameters l. 
Each instance of a HiP-SSM is an SSM conditioned on l. We also make the additional assumption that the parameter vector 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
 

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 z_t and 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 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 l based on an observed context set 𝒞_l. For any HiP-RSSM instance defined on a target sequence 𝒯=(o_t,a_t,o_t+1,...,o_t+N,a_t+N), over which we intend to perform state estimation/prediction, we maintain a fixed context set 𝒞_l. 𝒞_l in our HiP-SSM formalism can be obtained according to the algorithm designer's choice. We choose a 𝒞_l consisting of a set of tuples 𝒞_l = {o_t-n^l,a_t-n^l,o_t-n+1^l}_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 l∈ℒ based on an observed context set 𝒞_l = {x_n^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 𝒞_l is directly aggregated into the statistical description of l based on a factorized Gaussian observation model of the form p(r_n^l|l), where

    r_n^l    = enc_r(o_t-n^l,a_t-n^l,o_t-n+1^l), 
    σ_n^l    = enc_σ(o_t-n^l,a_t-n^l,o_t-n+1^l).

Here n is the index of an element from the context set 𝒞_l. 

Given a prior p_0(l) = 𝒩(l|μ_0,diag(σ_0)) we can compute the posterior p(l|𝒞_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(l|𝒞_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

    σ_l    = ( (σ_0)^⊖+ ∑_n=1^N (σ_n^l)^⊖)^⊖, 
    μ_l    = μ_0 + σ_l⊙∑_n=1^N (r^l_n - μ_0) ⊘σ_n^l

where ⊖, ⊙ and ⊘ denote element-wise inversion, product, and division, respectively. Intuitively, the mean of the latent task variable μ_l is a weighted sum of the individual latent observations r^l_n, while the variance of the latent task variable σ_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 

    p(z_t|w_1:t-1,a_1:t,𝒞_l) = ∬ p(z_t|z_t-1,a_t,l)   p(z_t-1|w_1:t-1,a_1:t-1,𝒞_l) 
       p(l|𝒞_l)dz_t-1dl.
 

We assume a dynamical model of the following form: p(z_t|z_t-1,a_t,l) = 𝒩(A_t-1z_t-1 + Ba_t + Cl,  Σ_trans) and
denote the posterior from the previous time-step by p(z_t-1|w_1:t-1,a_1:t-1,𝒞_l) = 𝒩(z_t-1^+, Σ_t-1^+). 

Following <cit.>, we assume that the action a_t is known and not subject to noise. 

At any time t, the posterior over the belief state z_t-1|w_1:t-1,a_1:t-1, posterior over the latent task variable l|𝒞_l and the action a_t are independent of each other since they form a “common effect" / v structure koller2007graphical with the unobserved variable 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(z_t|w_1:t-1,a_1:t,𝒞_l) = 𝒩(z_t^-,Σ_t^-), where

    z_t^-    =A_t-1z_t-1 + Ba_t + Cl, 
    Σ_t^-    = A_t-1Σ_t-1^+A_t-1^T + CΣ_lC^T + Σ_trans.





[Linear Combination Gaussian Marginalization]gigausslinearmarg
If u∼𝒩(μ_u + b,Σ_u) and v∼𝒩(μ_v,Σ_v)  are normally distributed independent random variables and if the conditional distribution p(y|u, v) = 𝒩(Au + b + Bv, Σ), then marginal p(y) = ∫ p(y|u,v)p(u)p(v)) dudv = 𝒩(Aμ_u + b + Bμ_v, AΣ_uA^T + BΣ_vB^T +  Σ).




The proof for the above identity is given in Appendix.





  
Modelling Choices We use a locally linear transition model 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 μ_l and σ_l directly into the latent space of the state-space model through additive interactions. The corresponding time update equations are given below:

    z_t^-    =A_t-1z_t-1^+ + b(a_t) + f(μ_l), 
    Σ_t^-    = A_t-1Σ_t-1^+A_t-1^T + f(σ_l)+ Σ_trans.

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 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 A and 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(z_t|o_1:t,a_1:t,C_l). We first map the observations at each time step o_t to a latent space using an observation encoder haarnoja2016backprop,becker2019recurrent that emits latent features w_t along with uncertainty in those features via a variance vector σ_obs^t. We then computed the posterior belief p(z_t|w_1:t,a_1:t,C_l), based on p(z_t|w_1:t-1,a_1:t,C_l) obtained from the time update, the latent observation (w_t,σ_t^obs) and the observation model H. This is exactly the Kalman update step, which has a closed form solution as shown below for a time instant t, 

!
  

    Kalman Gain:   Q_t    = Σ_t^- H^T(HΣ_t^- H^T  + I·σ_t^obs)^-1, 
       Posterior Mean:   z_t^+    = z_t^- + Q_t (w_t - Hz_t^-  ), 
       Posterior Covariance:   Σ_t^+    = (I - Q_t H) Σ_t^-,

  

where 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 𝒪(L^3), where L is the dimension of the latent state 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 𝒪(L^3) to 𝒪(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 z, given a prior p(z) and an observation model p(w|z), using the Bayes rule: p(z|w) = p(w|z)p(z)/p(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(w|z) = 𝒩(Hz, Σ_obs).



§ HIP-SSM AS ADAPTIVE MULTI-TASK WORLD MODELS





 §.§ 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>.

ruled




  
Batch Training. Let T  ∈ℛ^B × N × 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 ∈ T, in which the dynamics is modeled, we also input a set of previous interactions N, which forms the context set C ∈ℛ^B × N × 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 A_t, the nonlinear control factor b, the linear/nonlinear latent task transformation model C, the transition noise Σ_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., 

    L     =∑_t=1^H log p(o_t+1|w_1:t, a_1:t,  C_l) 
        = ∑_t=1^H log∫ p(o_t+1|z_t+1)  p(z_t+1|w_1:t,  a_1:t,  C_l) dz_t+1


The extension to multiple trajectories is straightforward and was omitted to keep the notation uncluttered. Here, 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(o_t+1|w_1:t, a_1:t,  C_l) ≈𝒩(μ_o_t+1,diag(σ_o_t+1)) where we use the mean of the prior belief μ_z_t+1^- to decode the predictive mean, that is, μ_o_t+1 =  dec_μ(μ_z_t+1^-) and the variance estimate of the prior belief to decode the observation variance, that is, σ_o_t+1 = dec_σ(Σ_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  

    ℒ(o_(1:T)) = 
       1T∑_t=1^T  log𝒩(o_t | dec_μ(z_t^+), dec_Σ (σ^u,+_t,
    σ^s,+_t,
    σ^l,+_t )),

where dec_μ(·) and dec_Σ(·) 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. 

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. 



ruled




§ 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.



 §.§ 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.



 §.§ 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.
r.45
0.5

    No Imputation     50% Imputation 

HiP-RSSM     2.96 ± 0.212     6.15 ± 0.327 

RKN     7.17±0.017     14.66 ± 0.224 

LSTM     9.14±0.026     51.21 ± 0.431 

GRU     9.216±0.073     53.14 ± 0.242 

FFNN     8.72±0.021     - 

NP     4.57±0.013     - 

MAML     5.04±0.051     - 



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


We finally visualize the latent task representations using TSNE in Figure <ref>. As seen in the plot, the HiP-SSM instances under similar terrain slopes cluster in the latent task space, indicating that the model correctly identifies the causal effects in the dynamics in an unsupervised fashion. 



 §.§ Ablation On Context Encoder For Inferring Task Abstraction

In table <ref>, we report the details of evaluating our Bayesian aggregation based context set encoder (discussed in <ref>) against a causal / recurrent encoder that takes into account the temporal structure of the context data owing to its sequential nature. We used a probabilistic recurrent encoder becker2019recurrent, whose mean and variance from the last time step is used to infer the posterior latent task distribution p(l|𝒞_l) = 𝒩(μ_l, diag(σ_l^2)). The dimensions of the latent task parameters obtained from both the set and recurrent encoders remain the same (60). 

The reported experiments are conducted on data from a wheeled mobile robot discussed in Section <ref>. As reported in Table <ref> the permutation invariant set encoder outperforms the recurrent encoder by a good margin in terms of prediction accuracy for both fully and partially observable scenarios. Additionally the set encoder is far more efficient in terms of computational time required for training as seen from the time taken per epoch for each of these cases as it allows for efficient parallelization.




 §.§ Visualizing Changing Hidden Parameters At Test Time Over Trajectories With Varying Dynamics

[11]r.35
 0.95
    < g r a p h i c s >

 An ablation on the performance of different task transformation models discussed in section <ref>. 
 
 
We finally perform inference using the trained HiP-RSSM in a multi-task / changing dynamics scenario where the dynamics continuously changes over time. We use the inference procedure described in appendix <ref> based on a fluid definition for “task” as the local dynamics in a temporal segment. We plot the global variation in the latent task variable captured by each instance of the HiP-RSSM over these local temporal segments using the dimensionality reduction technique UMAP mcinnes2018umap. As seen in figures <ref> and <ref>, the latent task variable captures these causal factors of variations in an interpretable manner.





§ CONCLUSION

We proposed HiP-RSSM, a probabilistically principled recurrent neural network architecture for modelling the world under non-stationary dynamics. We start by formalizing a new framework, HiP-SSM, to address the multi-task state-space modelling setting. HiP-SSM assumes a shared latent state and action space across tasks but additionally assumes latent structure in the dynamics.  We exploit the structure of the resulting Bayesian network to learn a universal dynamics model with latent parameter l via exact inference and backpropagation through time. The resulting recurrent neural network, namely HiP-RSSM, learns to cluster SSM instances with similar dynamics together in an unsupervised fashion. Our experimental results on various robotic benchmarks show that HiP-RSSMs significantly outperform state of the art recurrent neural network architectures on dynamics modelling tasks. We believe that modelling the dynamics in the latent space which disentangles the state, action and task representations can benefit multiple future applications including planning/control in the latent space and causal factor identification.


CHAPTER: MULTI TIME SCALE SSMS: HIERARCHICAL WORLD MODELS AT MULTIPLE TEMPORAL ABSTRACTIONS



This chapter is based on "Multi Time Scale World Models" shaj2023multi.


This chapter of the thesis explores the concept of "temporal depth" within generative models of the world by structuring them through nested hierarchies that span across different spatio-temporal scales. This approach addresses a significant limitation found in the prevailing literature on world models: the reliance on a singular temporal depth or time scale. Present formulations in the fields of machine learning and artificial intelligence typically operate with a single, finely detailed temporal scale, such as milliseconds. Singular time scale often limits accurate prediction and planning over long hoizons lecun2022path. For efficient long-term prediction and planning, the model must predict at multiple levels of temporal abstractions sutton1995td,precup1997multi,lecun2022path.


The world is rife with regularities. Day follows night, seasons follow one another, milk goes sour, faulty brakes are often followed by accidents, and so on. Regularities come at different time scales, ranging from tens of milliseconds to hundreds, to seconds, minutes, and upwards towards regularities or rules that are stable over weeks, months, and years.  For example, commuting to work involves planning at a higher level slow time scale (in minutes or hours) like which route to drive (or walk) and at a lower level fast time scale (in seconds or milliseconds) like how to brake/steer or choose precise skeletal muscle actions. Similarly, in robotic manipulation, the robot must be able to perform precise and coordinated movements to grasp and manipulate the object at a fast time scale while at a slower time scale the robot must also be able to recognize and utilize higher-level patterns and structures in the task, such as the shape, size, and location of objects, and the overall goal of the manipulation task. Generative models with "temporal depth" in the form of nested hierarchies offers several advantages in dealing with such scenarios lee2003hierarchical, hohwy2013predictive, lecun2022path as listed below:


    
  * Reflection of Causal Depth: The world is inherently structured hierarchically hohwy2013predictive, with causes and effects nested within one another on various scales of time and space (for example, in autonomous driving applications, slow regularities like traffic congestion trends including peak and off-peak hours, holiday traffic patterns, and seasonal changes in road conditions affect the fast regularities like real-time decisions on route navigation and speed adjustments). A hierarchical model allows the representation of this causal depth, enabling intelligent agents to understand not just immediate sensory inputs but also the broader and interconnected causal relationships that govern the external world.

    
  * Better Long Term Prediction: By organizing regularities in the world from faster, detailed levels to slower, more abstract levels, hierarchical models enable predictions at multiple scales of precision and time. This flexibility is fundamental to adaptive behavior, allowing for both immediate, detailed responses and longer-term planning based on abstracted patterns of regularity friston2008hierarchical,hohwy2013predictive. Fast changing regularities are good for detail; slower regularities are more general and abstract. This makes sense when we consider what regularities allow us to predict. If I want to predict something with great perceptual precision, then I cannot do it very far into the future, so I need to rely on a fast changing regularity. On the other hand, predictions further into the future come at a loss of precision and often detail.

    
  * Adaptability and Transferability: As seen in Chapter <ref> hierarchical temporal abstractions can capture relevant task structures across dynamical systems under non-stationary conditions, which can be used to identify similarities and differences between tasks. Learning dynamics of these nested hierarchical abstractions allow for making "predictions about predictions" and subsequent transfer of knowledge learned from one task to another shanahan2022abstraction, lecun2022path. This adaptability is key to navigating a world that is constantly changing, allowing for the refinement of predictions and behaviors in response to new information.
  
 



This chapter of the dissertation attempt to come up with a principled probabilistic formalism for learning such multi-time scale world models as a hierarchical sequential latent variable model. We show that such models can better capture the complex, nonlinear dynamics of a system more efficiently and robustly than models that learn on a single timescale. This is exemplified in several challenging simulated and real-world prediction tasks such as the D4RL dataset, a simulated mobile robot, and real manipulators including data from heavy machinery excavators. 



§ RELATED WORK


Multi Time Scale World Models One of the early works that enabled environment models at different temporal scales to be intermixed, producing temporally abstract world models was proposed by <cit.>. The work was limited to tabular settings but showed the importance of learning environment dynamics at multiple abstractions. However, there have been limited works that actually solve this problem at scale as discussed in <cit.>. A probabilistically principled formalism for these has been lacking in literature and this work is an early attempt to address this issue.
Deep State Space Models. Deep SSMs combine the benefits of deep neural nets and SSMs by offering tractable probabilistic inference and scalability to high-dimensional and large datasets. <cit.> use neural network architectures based on exact inference on SSMs and perform state estimation and dynamics prediction tasks. <cit.> extend these models to modelling non-stationary dynamics. <cit.> perform learning and inference on SSMs using variational approximations. However, most of these recurrent state-space models have been evaluated on very short-term prediction tasks in the range of a few milliseconds and model the dynamics at a single time scale.
Transformers 
Recent advancements in Transformers vaswani2017attention, radford2019language, brown2020language, which rely on attention mechanism, have demonstrated superior performance in capturing long-range dependency compared to RNN models in several domains including time series forecasting zhou2021informer,liu2022nonstationary and learning world models micheli2023transformers. <cit.> use transformer architectures based on a direct multistep loss zeng2022transformers and show promising results for long-term forecasting since they avoid the accumulation of errors from autoregression. On the other hand <cit.> uses a GPT-like autoregressive version of transformers to learn world models. These deterministic models, however, do not deal with temporal abstractions and uncertainty estimation in a principled manner. Nevertheless, we think Transformers that operate at multiple timescales based on our formalism can be a promising alternative research direction.


§ MULTI TIME SCALE STATE SPACE MODELS


Our goal is to learn a principled sequential latent variable model that can model the dynamics of artificial embodied agents (specifically robotic systems) under multiple levels of temporal abstractions based on observed sensory data. To do so, we introduce a new formalism, called Multi-Time Scale State Space (MTS3) Model, with the following desiderata:
i) It is capable of modeling dynamics at multiple time scales. ii) It allows for a single global model to be learned that can be shared across changing configurations of the environments. iii) It can provide accurate long-term predictions and uncertainty estimates. iv) It is probabilistically principled, yet scalable during learning and inference.

We start by introducing a formal definition of a 2 time scale Multi-time scale SSM in Section <ref> and then proceed to a general definition of an MTS3 with an arbitrary number of hierarchies/timescales in Section <ref>.



 §.§ Formal Definition Of a 2-Level MTS3


An MTS3 model with 2 timescales is defined by two SSMs on a fast and a slow time scale respectively. Both SSMs are coupled via the latent state of the slow time scale SSM, which parameterizes / “reconfigures” the system dynamics of the fast time scale SSM. While the fast time scale SSM runs at the original time step Δ t of the dynamical system, the slow time scale SSM is only updated every H step, i.e., the slow time scale time step is given by H Δ t. We will derive closed-form Gaussian inference for obtaining the beliefs for both time scales, resulting in variations of the Kalman update rule which are also fully differentiable and used to back-propagate the error signal becker2019recurrent,haarnoja2016backprop. The definition with a 2-level MTS3 along with the inference and learning schemes that we propose is directly extendable to an arbitrary number of temporal abstractions by introducing additional feudal dayan1992feudal hierarchies with longer discretization steps and is further detailed in Section <ref>.



  §.§.§ Fast time-scale SSM


  The fast time-scale (fts) SSM is given by 𝒮_fast = ( 𝒵, 𝒜, 𝒪, f_ l^fts, h^fts, Δ t, ℒ). Here, l ∈ℒ is a task descriptor that parameterizes the dynamics model of the SSM and is held constant for H steps. We will denote the task descriptor for the kth time window of H steps as l_k. The probabilistic dynamics and observation model of the fast time scale for the tth time step in the kth window can then be described as  
 
    p(z_k,t|z_k,t-1,a_k,t-1,  l_k)     = 𝒩(f^fts_ l(z_k,t-1,a_k,t-1,  l_k), Q),
        and 
    
       p(o_k,t|z_k,t)     = 𝒩(h^fts(z_k,t), R).

 

  
Task-conditioned marginal transition model.
Moreover, we have to consider the uncertainty in the task descriptor (which will, in the end, be estimated by the slow time scale model), i.e., instead of considering a single task descriptor l_k, we have to consider a distribution over task descriptors p( l_k) for inference in the fts-SSM. This distribution will be provided by the slow-time scale SSM for every time window k. We can further define the marginal task-conditioned transition model for the time window k that is given by 

    p_ l_k(z_k,t|z_k,t-1,a_k,t-1)     = ∫ p(z_k,t|z_k,t-1,a_k,t-1,  l_k) p( l_k) d  l_k



  
Latent observations.
Following <cit.>, we replace the observations by latent observations and their uncertainty, i.e., we use latent observation encoders to obtain w_k,t = enc_w( o_k,t) and an uncertainty encoder σ_k,t = enc_σ( o_k,t). The observation model is thus given by p( w_k,t| z_k,t) = 𝒩(h^fts( z_k,t), diag(σ_k,t)).



  §.§.§ Slow time-scale SSM


The slow time scale (sts) SSM only updates every time step H and uses the task parameter l as a latent state representation. Formally, the SSM is defined as 𝒮_slow = ( ℒ, ℰ, 𝒯, f^sts, h^sts, H Δ t). It uses an abstract observation β∈ℬ and abstract action α∈𝒜 that summarize the observations and actions, respectively, throughout the current time window. The general dynamics model is hence given by 
    p( l_k| l_k-1, α_k) = 𝒩(f^sts( l_k-1, α_k),  S).


Although there exist many ways to implement the abstraction of observations and actions of time windows, we choose to use a consistent formulation by fusing the information from all H time steps of time window k using Gaussian conditioning. 
 


  
Observation abstraction.
In terms of the abstract observation model, we choose to model H observations β_k,t, t ∈ [1,H] for a single slow-scale time step k. All these observations can then be straightforwardly integrated into the belief state representation using incremental observation updates. The abstract observation and its uncertainty for time step t is again obtained by an encoder architecture, i.e, 

    β_k,t = enc_β( o_k,t, t),   ν_k,t = enc_ν( o_k,t, t),

and p(β_k,t| l_k) = 𝒩(h^sts( l_k), diag(ν_k,t)).
Hence, the abstract observation β_k,t contains the actual observation o_k,t at time step t as well as a temporal encoding for the time-step. Although multiple Bayesian observation updates are permutation invariant, the temporal encoding preserves the relative time information between the observations, similar to current transformer architectures.



  
Action abstraction. 
The abstract action α_k causes transitions to latent task l_k from l_k-1. It should contain relevant information of all primitive actions a_k,t, t ∈ [1,H] executed in the time window k. To do so, we again use Bayesian conditioning and latent action encoding. Each control action a_k,t and the encoding of the time step t are encoded in its latent representation and its uncertainty estimate, i.e.,

    α_k,t = enc_α( a_k,t, t),   ρ_k,t = enc_ρ( a_k,t, t).

[13]r.31
[b]0.31
   width=
         
     
Generative model for the abstract action α_k. The hollow arrows are deterministic transformations leading to implicit distribution α_k,t using an action set encoder.


Single latent actions α_k,t can be aggregated into a consistent representation α_k using Bayesian aggregation <cit.>. To do so, we use the likelihood p(α_k,t|α_k) = 𝒩(α_k, diag( ρ_k,t)) and obtain the posterior p(α_k|α_k,1:H) = 𝒩(μ_α_k, Σ_α_k), which is obtained by following the standard Bayesian aggregation equations, see Appendix <ref>. Note that our abstract action representation also contains an uncertainty estimate which can be used to express different effects of the actions on the uncertainty of the prediction. Due to the Gaussian representations, we can compute the marginal transition model 

    p_α_k( l_k| l_k-1, α_k,1:H) = ∫ p_α_k( l_k| l_k-1, α_k) p(α_k |α_k,1:H) dα_k.

This transition model is used for inference and its parameters are learned. 



  §.§.§ Connecting both SSMs via inference

In the upcoming sections, we will devise Bayesian update rules to obtain the prior p( l_k| β_1:k-1,  α_1:k) and posterior p( l_k| β_1:k,  α_1:k) belief state for the sts-SSM as well as the belief states for the fts-SSM. The prior belief p( l_k| β_1:k-1,  α_1:k) contains all information up to time window k-1 and serves as a distribution over the task-descriptor of the fts-SSM, which connects both SSMs. This connection allows us to learn both SSMs jointly in an end-to-end manner.  

The probabilistic graphical model of our MTS3 model is depicted in Figure <ref>. In the next section, we will present the detailed realization of each SSM to perform closed-form Gaussian inference and end-to-end learning on both time scales.



 §.§ Inference in the Fast Time-Scale SSM


The fts-SSM performs inference for a given time window k of horizon length H. To keep the notation uncluttered, we will also omit the time-window index k whenever the context is clear. 
We use a linear Gaussian task conditional transition model, i.e, 
    p(z_t|z_t-1,a_t-1,l_k) = 𝒩(Az_t-1 + Ba_t-1 + Cl_k,   Q),
 where A, B, C and Q are state-independent but learnable parameters.
In our formulation, the task descriptor can only linearly modify the dynamics which was sufficient to obtain state-of-the-art performance in our experiments, but more complex parametrizations, such as locally linear models, would also be feasible. Following <cit.>, we split the latent state z_t = [ p_t,  m_t]^T into its observable part p_t and a part m_t that needs to be observed over time. 
We also use a linear observation model p( w_t| z_t) = 𝒩( H  z_t, diag(σ_t)) with H = [ I,  0]. 

We will assume that the distribution over the task descriptor is also given by a Gaussian distribution, i.e., p( l_k) = 𝒩(μ_ l_k, Σ_ l_k), which will be provided by the slow-time scale (sts) SSM, see Section <ref>. Given these modelling assumptions, the  task variable can now be integrated out in closed form, resulting in the following task-conditioned marginal transition model

    p_ l_k(z_t|z_t-1,a_t-1) 
     = 𝒩(Az_t-1 + Ba_t-1 + Cμ_ l_k,  Q +  C Σ_ l_k C^T ),

which will be used instead of the standard dynamics equations.

 We follow the same factorization assumptions as in <cit.> and only estimate the diagonal elements of the block matrices of the covariance matrix of the belief, see Appendix B. The update equations for the Kalman prediction and observation updates are therefore equivalent to the RKN becker2019recurrent.




 §.§ Inference in the Slow-Time Scale SSM




  
Prediction Update.
We follow the same Gaussian inference scheme as for the fts-SSM, i.e., we again employ a linear dynamics model 
p( l_k| l_k-1, α_k) = 𝒩( X  l_k-1 +  Y α_k,  S), where X, Y and S are learnable parameters. The marginalized transition model for the abstract actions is then given by 


    p_α_k(l_k|l_k-1)     = ∫ p(l_k|l_k-1,α_k) p(α_k) d α_k 
     = 𝒩(Xl_k-1 + Yμ_α_k, S +  Y Σ_α_k Y^T ).


We can directly use this transition model to obtain the Kalman prediction update which computes the prior belief p_α_1:k( l_k | β_1:k-1) = 𝒩(μ_l_k^-, Σ_l_k^-) from the posterior belief  p_α_1:k-1( l_k-1 | β_1:k-1) = 𝒩(μ_l_k-1^+, Σ_l_k-1^+) of the previous time window, see Appendix <ref>. 



  
Observation/Task Update.
Similarly, we will use a linear observation model for the abstract observations 
p(β_k,t|  l_k) = 𝒩( H  l_k, diag(ν_k,t)) with H = [ I,  0].
As can be seen from the definition of the observation matrix H, the latent space is also decomposed into its observable and unobservable part, i.e., l_k = [ u_k,  v_k]. In difference to the standard factorized Kalman observation update given in Section <ref>, we have to infer with a set of observations β⃗_k,t with t = 1 … H for a single time window k. While in principle, the Kalman observation update can be applied incrementally H times to obtain the posterior p_α_1:k( l_k | β_1:k) = 𝒩(μ_l_k^+, Σ_l_k^+), such an update would be very slow and also cause numerical inaccuracies.[13]r0.55 
    
    
    
    < g r a p h i c s >

  Implementation of task update layer which performs posterior latent task inference in the sts-SSM.
  
 Hence, we devise a new permutation invariant version of the update rule that allows parallel processing with set encoders zaheer2017deep. We found that this update rule is easier to formalize using precision matrices. Hence, we first transform the prior covariance vectors σ_l_k^u-, σ_l_k^l- and σ_l_k^s- to its corresponding precision representation λ_l_k^u-, λ_l_k^l- and λ_l_k^s- which can be performed using block-wise matrix inversions of Σ_l_k^-. Due to the factorization of the covariance matrix, this operation can be performed solely by scalar inversions. 

As the update equations are rather lengthy, they are given in Appendix <ref>, <ref>. Subsequently, we compute the posterior precision, where only λ_l_k^u is changed by 

    λ_l_k^u+ = λ_l_k^u- + ∑_t=1^H  1 ⊘ν_k,t

 while λ_l_k^l+ = λ_l_k^l- and λ_l_k^s+ = λ_l_k^s- remain constant. The operator ⊘ denotes the element-wise division. From the posterior precision, we can again obtain the posterior covariance vectors σ_l_k^u+, σ_l_k^l+ and σ_l_k^s+ using only scalar inversions, see Appendix <ref>, <ref>. The posterior mean μ_l,k^+ can now be obtained from the prior mean μ_l,k^- as
 
    [c]
       μ_l,k^+=μ_l,k^- + [ [ σ_l_k^u+; σ_l_k^s+;          ]]  ⊙[ [ ∑_t=1^H(β_k,t-μ^u,-_l_k) ⊘ν_k,t; ∑_t=1^H(β_k,t-μ^u,-_l_k) ⊘ν_k,t;                                 ]].

For H = 1, i.e a single observation, the given equation is equivalent to the factorized Kalman observation update becker2019recurrent. Furthermore, the given rule constitutes a unification of the batch update rule for Bayesian aggregation volpp2020bayesian and the incremental Kalman update for our factorization of belief state representation becker2019recurrent detailed in Section <ref> and Appendix <ref>. 



 §.§ A General Definition For an  N-level MTS3
 

The human brain's internal representation of the real world is complex, comprising more than just two distinct causal hierarchies. Neuroscience research hohwy2013predictive,friston2008hierarchical,jiang2021predictive suggests that these systems exhibit multiple hierarchical structures. Consequently, we propose a formal definition for Multi-Hierarchy Theory of Mind (MTS3) that accommodates an arbitrary number (N) of hierarchies.

An N-level MTS3 can be defined as a family of N-state space models, {S_0, S_1, ..., S_N-1}. Each of the state space models S_i is given by S_i = (Z_i, A_i, O_i, f_i, h_i, H_i Δ t, L_i), where Z_i is the state space, A_i the action space, and O_i the observation space of the SSM. The parameter H_i Δ t denotes the discretization time step and f_i and h_i the dynamics and observation models, respectively. Here, l_i ∈ L_i is a task descriptor that parameterizes the dynamics model of the SSM and is constant for a local window of steps H_i+1. l_i is a function of the latent state of the SSM one level above it, i.e., S_i+1. The boundary cases can be defined as follows: for i=0, H_0 = 1. Similarly, for i=N-1, the latent task descriptor L_i is an empty set. For all i, H_i < H_i+1.




 §.§ Inference In N-Level MTS3

As seen from the definition, an N-level MTS3 is a set of N SSMs that are strictly feudal dayan1992feudal from top to bottom. Top-level SSM (managers) makes decisions/predictions independently of the bottom-level SSMs, while the bottom-level SSM (worker) is conditioned on its immediate manager. Each of these SSMs performs inference via two stages at every time-step, i.e, 
    
  * Gaussian Conditioning Updating the posterior belief based on incoming observations.
    
  * Gaussian Marginalization Estimating the prior belief for the next step via learned linear dynamics. Note that this stage is conditioned on the current belief of an immediate manager if present.
 

The efficient Gaussian conditioning and marginalization processes for State-Space Models (SSMs) across all levels in the MTS3 model are specific instances of two broad Gaussian principles derived within the thesis. These principles are formally presented as Gaussian Identity <ref> and Gaussian Identity <ref>, as detailed below:

[Gaussian Conditioning]gigausscond
     
    Consider the graphical model given in Figure <ref>, where a set of N conditionally i.i.d observations r = {r_i}_i=1^N are generated by a latent variable l and the observation model p(r_i|l) = 𝒩(r_i |Hl, diag(σ_i^obs)). Assuming an observation model H=[I,0], the mean (μ) and precision matrix (Λ) of the posterior over the latent variable l, p(l|r) = 𝒩( μ_l^+, Σ_l^+) = 𝒩( μ_l^+, (Λ_l^+)^-1), given the prior p_0(l) = 𝒩( μ_l^-, Σ_l^-) = 𝒩( μ_l^-, (Λ_l^-)^-1) have the following permutation invariant closed form updates.

    Λ_l^+    = Λ_l^- + [ [ diag(∑_i=1^n1/σ_i^obs),                       0;                     0 ,                       0;                         ]] 
    μ_l^+   =μ_l^- + [ [ σ_l^u+; σ_l^s+;        ]]  ⊙[ [ ∑_i=1^N(r_i-μ^u,-_l) ⊙1/σ_i^obs; ∑_i=1^N(r_i-μ^u,-_l) ⊙1/σ_i^obs;                                 ]]



The proof for the above identity is given in Appendix <ref>.




The closed form updates for the resulting posterior distribution p(l|r) is permutation invariant with respect to the observation set r.

Corollary <ref> is also a direct consequence of Theorem 1, which states that the sequential Bayesian update is permutation invariant. The permutation invariance/exchangeability property has both computational and theoretical implications in efficient posterior inference over latent task variables, which is further discussed in the section.

[Permutation Invariance Of Bayesian Inversion]thmperminv

For any conditionally i.i.d model where you have a global parameter θ, and a set of observations X= {x_i}_i=1^N drawn conditionally i.i.d from a distribution p(X|θ), then for any permutation π,  the posterior p(θ| x_1, …, x_N) = p(θ|π(x_1), …, π(x_N)). Thus the posterior p(θ| X) is permutation invariant with respect to the set X. 




The proof of the above theorem can be found in Appendix <ref>.


The derived update equations can be coded as a layer in the neural network architecture as shown in Figure <ref>.





[Linear Combination Gaussian Marginalization]gigaussmarg

Consider the graphical model in Figure <ref>, where a set of N normally distributed independent random variables u = {u_i∼𝒩(μ_u_i,Σ_u_i)}_i=0^N forms a “common effect/V Structure” with a latent variable y. If the conditional distribution p(y|u_1, u_2, . ., u_N) = 𝒩( ∑_i=0^N A_iu_i, Σ), then marginal p(y) = ∫ p(y|u_1, u_2, . ., u_N)∏_i=1^Np(u_i)du_i = 𝒩(∑_i=1^N A_iμ_u_i , Σ + ∑_i=1^N A_iΣ_u_iA_i^T).




The proof for the above identity is given in Appendix <ref>.



Even though our experiments focus on MTS3 models with 2 hierarchies, extensive experimentation with more hierarchies can be taken as future work.



§ MTS3 AS A HIERARCHICAL WORLD MODEL

 

MTS3 allows for a natural way to build world models that can deal with partial observability, nonstationarity, and uncertainty in long-term predictions, properties which are critical for model-based control and planning. Furthermore, introducing several levels of latent variables, each working at a  different time scale allows us to learn world models that can make action conditional predictions/“dreams” at multiple time scales and multiple levels of state and action abstractions. 

 




 §.§ Conditional Multi Time Predictions With World Model
 

Conditional multistep ahead predictions involve estimating plausible future states of the world resulting from a sequence of actions. Our principled formalism allows for action-conditional future predictions at multiple levels of temporal abstraction. The prediction update for the sts-SSM makes prior estimates of future latent task variables conditioned on the abstract action representations. However, the task conditional prediction update in the fts-SSM estimates the future prior latent states, conditioned on primitive actions and the inferred latent task priors, which are decoded to reconstruct future observations. 
To initialize the prior belief p( z_k,1) for the first time step of the time window k, we use the prior belief p( z_k-1, H+1) for the last time step of the time window k-1.








 §.§ Optimizing the Predictive Log-Likelihood
 
The training objective for the MTS3 involves maximizing the posterior predictive log-likelihood which is given below for a single trajectory, i.e., 

    L     = ∑_k=1^N ∑_t=1^H log p(o_k,t+1|β_1:k-1,α_1:k,w_k,1:t, a_k,1:t) 
        
         = ∑_k=1^N ∑_t=1^H log∬ p(o_k,t+1|z_k,t+1)  p(z_k,t+1|w_k,1:t,  a_k,1:t,  l_k) 
       p( l_k|β_1:k-1,α_1:k)dz_k,t+1 d  l_k 
        =  ∑_k=1^N ∑_t=1^H log∫ p(o_k,t+1|z_k,t+1)  p_ l_k(z_k,t+1|w_k,1:t,  a_k,1:t) dz_k,t+1.

The extension to multiple trajectories is straightforward and was omitted to keep the notation uncluttered. Here, o_k,t+1 is the ground-truth observations at time step t+1 and time window k that needs to be predicted from all (latent and abstract) observations up to time step t. The prior belief corresponding to the latent state p_ l_k(z_k,t+1|w_k,1:t,  a_k,1:t) has a closed form solution as discussed in Section <ref>. 


We employ a Gaussian approximation of the posterior predictive log-likelihood of the form p(o_k,t+1|β_1:k-1,α_1:k,w_k,1:t, a_k,1:t) ≈𝒩(μ_o_k,t+1,diag(σ_o_k,t+1)) where we use the mean of the prior belief μ_z_k,t+1^- to decode the predictive mean, i.e, μ_o_k,t+1 =  dec_μ(μ_z_k,t+1^-) and the variance estimate of the prior belief to decode the observation variance, i.e., σ_o_k,t+1 = dec_σ(Σ_z_k,t+1^-). This approximation can be motivated by a moment matching perspective and allows for end-to-end optimization of the logarithmic likelihood without using auxiliary objectives such as the ELBO <cit.>.

Gradients are computed using (truncated) backpropagation over time (BPTT) werbos1990backpropagation and clipped.  
We optimize the objective using the Adam kingma2014adam stochastic gradient descent optimizer with default parameters. We refer to Appendix A for more details. For training, we also initialize the prior belief p( z_k,1) with the prior belief p_ l_k-1( z_k-1,H+1| w_k-1,1:H, a_k-1,1:H) from the previous time window k-1. However, we cut the gradients for the fast time scale between time windows as this avoids vanishing gradients, and we observed a more stable learning behavior. Yet, the gradients can still flow between time windows for the fts-SSM via the sts-SSM.
 


 §.§ Imputation Based Self Supervised Training For Long Term Prediction 
 

Using the given training loss 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 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. This imputation scheme is applied at both time scales, masking out single time steps or whole time windows of length H. The imputation mask is also randomly resampled for every minibatch.





§ EXPERIMENTS

In this section, we evaluate our approach to a diverse set of simulated and real-world dynamical systems for long-horizon prediction tasks. Our experiments are designed to answer the following questions. (a) Can MTS3 make accurate long-term deterministic predictions (mean estimates)? (b) Can MTS3 make accurate long-term probabilistic predictions (variance estimates)? (c) How important are the modelling assumptions and training scheme? 



 §.§ Baseline Dynamics Models

While a full description of our baselines can be found in Appendix <ref>, a brief description of them is given here: (a) RNNs - We compare our method to two widely used recurrent neural network architectures, LSTMs lstm and GRUs gru. 
(b) RSSMs - Among several RSSMs from the literature, we chose RKN becker2019recurrent and HiP-RSSM shaj2022hidden as these have shown excellent performance for dynamics learning for short-term predictions and rely on exact inference as in our case. (c) Transformers - We also compare with two state-of-the-art Transformer vaswani2017attention variants. The first variant (AR-Transformer) relies on a GPT-like autoregressive prediction radford2019language,brown2020language. Whereas the second variant (Multi-Transformer) uses direct multi-step loss zeng2022transformers from recent literature on long horizon time-series forecasting zhou2021informer,liu2022nonstationary,nie2023a. Here, multistep ahead predictions are performed using a single shot given the action sequences.









 §.§ Environments and Datasets


We experiment with three broad datasets. While full descriptions of these datasets, dataset creation procedure, and overall statistics are given in Appendix E, a brief description of them is as follows.
(a) D4RL Datasets - We use a set of 3 different environments/agents from D4RL dataset fu2020d4rl, which includes the HalfCheetah, Medium Maze and Franka Kitchen environment. Each of these was chosen because of their distinct properties like sub-optimal trajectories (HalfCheetah), realistic domains / human demonstrations (Kitchen), multi-task trajectories, non-markovian collection policies (Kitchen and Maze) and availability of long horizon episodes (all three). (b) Manipulation Datasets - We use 2 datasets collected from a real excavator arm and a Panda robot. The highly non-linear non-markovian dynamics due to hydraulic actuators in the former and non-stationary dynamics owing to different payloads in the latter make them challenging benchmarks. Furthermore, accurate modelling of the dynamics of these complex systems is important since learning control policies for automation directly on large excavators is economically infeasible and potentially hazardous. (c) Mobile Robotics Dataset - We set up a simulated four-wheeled mobile robot traversing a highly uneven terrain of varying steepness generated by a mix of sinusoidal functions. This problem is challenging due to the highly non-linear dynamics involving wheel-terrain interactions and non-stationary dynamics introduced by varying steepness levels. In all datasets, we only use information about agent/object positions and we mask out velocities to create a partially observable setting.





 §.§ Can MTS3 make accurate long-term deterministic predictions (mean estimates)?


Here we evaluate the quality of the mean estimates for long-term prediction using our approach. The results are reported in terms of "sliding window RMSE" in Figure <ref>. We see that MTS3 gives consistently good long-term action conditional future predictions on all 6 datasets. Deep Kalman models becker2019recurrent,shaj2022hidden which operate on a single time scale fail to give meaningful mean estimates beyond a few milliseconds. Similarly, widely used RNN baselines lstm,gru which form the backbone of several world models ha2018world,hafner2019learning give poor action conditional predictions over long horizons. AR-Transformers also fail possibly due to error accumulation caused by the autoregression. However, Multi-Transformers are a strong baseline that outperforms MTS3 in the Medium Maze and Panda dataset by a small margin. However, on more complex tasks like the Kitchen task, which requires modelling multi-object, multi-task interactions gupta2019relay, MTS3 is the only model that gives meaningful long horizon predictions. A detailed description of the metric "sliding window RMSE" is given in Appendix <ref>. A visualization of the predicted trajectories vs. ground truth is given in Appendix <ref>.




 §.§ Can MTS3 make accurate long-term probabilistic predictions (variance estimates)?

Next, we examine the question of whether the principled probabilistic inference translates to accurate uncertainty quantification during long-horizon predictions. We trained all the baselines with a negative
log-likelihood loss and used the same as a metric to quantify the quality of uncertainty estimates. During the evaluation we rely on a sliding window approach (see Appendix <ref>) and report the results for the last timestep in Table <ref>. As seen in Table <ref>, MTS3 gives the most accurate uncertainty estimates in all datasets except Medium Maze, where it is outperformed by Multi-Transformer. Also, notably, AR-Transformers and deep Kalman models fail to learn any meaningful uncertainty representation when it comes to long-term predictions. 





 §.§ How important are the modelling assumptions and training scheme?


Now, we look at three important modelling and training design choices: (i) splitting the latent states to include an unobservable “memory” part using observation model h^sts=h^fts= H = [ I,  0] as discussed in Sections <ref> and <ref>, (ii) action abstractions in Section <ref>, (iii) training by imputation. 
 


To analyze the importance of the memory component, we derived and implemented an MTS3 variant with an observation model of h^sts=h^fts= I and a pure diagonal matrix representation for the covariance matrices. As seen in Figure <ref>, this results in worse long-term predictions, suggesting that splitting the latent states in its observable and unobservable part in MTS3 is critical for learning models of non-markovian dynamical systems. Regarding (ii), we further devised another variant where MTS3 only had access to observations, primitive actions and observation abstractions, but no action abstractions. As seen in our ablation studies, using the action abstraction is crucial for long-horizon predictions.

Our final ablation (iii) shows the importance of an imputation-based training scheme discussed in Section <ref>. As seen in Figure <ref> when trained for 1 step ahead predictions without imputation, MTS3 performs significantly worse for long-term prediction suggesting the importance of this training regime.













 §.§ What is the role of the discretization step H.Δ t?


Finally, we perform ablation for different values of H.Δ t, which controls the time scale of the task dynamics. The results reported are for the hydraulics dataset. The higher the value of H, the slower the timescale of the task dynamics relative to the state dynamics. As seen in Figure Figure <ref>, smaller values of H (2,3,5 and 10) give significantly worse performance. Very large values of H (like 75) also result in degradation of performance. To further get an intuitive understanding of the MTS3's behaviour under different timescales, we plot the predictions given by MTS3 for different values of H on a trajectory handpicked from the hydraulics excavator dataset. As seen in Figure <ref> for large values of H like 30 and 75, we notice that the slow-changing task dynamics "reconfigures" the fast dynamics every 30 and 75-step window respectively, by conditioning the lower level dynamics with the newly updated task prior. This effect is noticeable as periodic jumps or discontinuities in the predictions, occurring at 30 and 75-step intervals. Also, for a very large H like 75, the fast time scale ssm has to make many more steps in a longer window resulting in error accumulation and poor predictions. A similar trend was observed for the mobile robot dataset. As seen in Figure <ref>, smaller values of H (like 2,3,5 and 10) give significantly worse performance. Very large values of H (like 150) also result in degradation of performance. In the paper, we used a value of H=75.





§ CONCLUSION AND FUTURE WORK

In this work, we introduce MTS3, a probabilistic formalism for learning the dynamics of complex environments at multiple time scales. By modelling the dynamics of the world at multiple levels of temporal abstraction we capture both the slow-changing long-term trends and fast-changing short-term trends in data, leading to highly accurate predictions spanning several seconds into the future. Our experiments demonstrate that simple linear models with principled modelling assumptions can compete with large transformer model variants that require several times more parameters. Furthermore, our inference scheme also allows for principled uncertainty propagation over long horizons across multiple time scales which capture the stochastic nature of environments. We believe our formalism can benefit multiple future applications including hierarchical planning/control. We discuss the limitations and broader impacts of our work in Chapter <ref>.


CHAPTER: OUTLOOK


This thesis questioned the current formalisms for learning-world models due to their inability to capture some critical criteria for a foundational world model. Specifically, we looked at three aspects (i) scalable probabilistic modelling that captures causal relations in our world (Chapter <ref>, <ref> and <ref>), (ii) adaptability to changing tasks (Chapter <ref>) and (iii) hierarchical modelling at multiple time scales and temporal abstractions (Chapter <ref>).
In this context the thesis proposed two new formalisms Hidden Parameter State Space Models (HiP-SSM in Chapter <ref>) and Multi Time Scale State Space Model (MTS3 in Chapter <ref>) besides making an existing formalism of SSMs more robust for performing interventions/counterfactuals with action/control signals (Ac-SSM in Chapter <ref>).

The proposed formalisms are in line with related theory in computational neuroscience called predictive processing and Bayesian brain hypothesis (Chapter <ref>). In the proposed generative models, Bayesian inversion is employed to deduce the underlying causes from the incoming sensory data at various levels of abstraction. The Bayesian inversion results in update rules that use "precision weighting" friston2009free,hohwy2013predictive,seth2014cybernetic as a mechanism of directing "attention" to the relevant sensory information. The idea of "precision estimation" by learned sensory encoders is also in line with discussions in the broader neuroscience community. The world model formalisms proposed are adaptable to changing situations (embodying cognitive flexibility) and can make top-down long-horizon predictions based on nested causal hierarchies. 

Machines that can replicate human intelligence and type 2 reasoning capabilities should be able to reason at multiple
levels of temporal abstraction using internal world models lecun2022path,gupta2024essential. The work represents a hopeful step towards developing embodied AI systems with a deeper, more intuitive grasp of the world, aligning machine learning processes more closely with natural intelligence. We believe that the findings of this dissertation will contribute to ongoing research in the creation of robust, principled, and scalable world models, while acknowledging the vast scope for further research and development in this area.



  
Limitations and Future Work 
The thesis specifically concentrated on the aspect of prediction rather than control or planning. Addressing hierarchical planning, much like hierarchical prediction, remains a significantly underexplored challenge. A logical progression for future research would be to extend the concept of planning within the exact inference framework botvinick2012planning,watson2020stochastic. This approach necessitates the derivation of backward messages and the use of the Kalman duality principle within the learned latent spaces of a deep encoder. Furthermore, investigating variational approaches to the hierarchical models presented, and considering hierarchical planning/control as a problem of approximate inference toussaint2009probabilistic,friston2009free,millidge2020relationship, presents another intriguing avenue for research.

A further limitation noted in the thesis is its reliance on proprioceptive sensors for experimentation, without empirical validation using high-dimensional sensory information, such as vision. An important avenue for future research would involve evaluating the usefulness and characteristics of abstractions derived through Bayesian inversion/aggregation when applied to high-dimensional and  multimodal sensory inputs. Conducting experiments with image-based data and exploring "non-reconstruction" based loss functions, as suggested by <cit.>, represent promising directions for upcoming work.

Though our hierarchical formalism outperform transformer variants on long horizon predictions on several tasks, a limitation for these to be used as foundational world models is the computational bottleneck due to the sequential nature of dynamics. However, there have been exciting progress recently in deterministic linear SSMs gu2021efficiently,smith2022simplified,mondal2023efficient that allow efficient parallelization during training. Since our dynamics is linear similar to these approaches, parallelization using similar techniques as employed by deterministic SSMs can be another important direction of future research.

Our experimentation was limited to the hierarchical models with just two levels of abstraction, as this configuration proved to be adequate for numerous tasks we carried out. Nevertheless, there is room for deeper exploration.The intersection of machine learning with broader cognitive theories, particularly the concept of machine consciousness, presents a novel yet underexplored frontier in the field. Traditionally, consciousness has been a topic more familiar to the realms of philosophy and psychology than to machine learning and computer science. This often makes it a delicate subject within our discipline.

As I write this thesis in the Spring of 2024, the machine learning field finds itself at a crossroads, significantly influenced by the impressive capabilities of attention/transformer vaswani2017attention based large language models (LLMs). These advancements suggests a pressing need for mathematical models or analytical frameworks to better understand machine consciousness bengio2017consciousness,chalmers2023could. Numerous theories on consciousness highlight the critical role of recurrent processing chalmers2023could,lamme2010neuroscience, a feature notably absent in current LLM architectures. In this context, I believe that hierarchical generative models, particularly those extending beyond the 2-level temporal depth as explored in this thesis, represent a promising research direction for embedding consciousness priors friston2013consciousness,bengio2017consciousness into machine learning systems. This approach could offer profound insights into the mechanisms underpinning consciousness and its potential replication in artificial systems.












[heading=bibintoc]






CHAPTER: APPENDIX: ACTION CONDITIONAL SSM



§ INVERSE DYNAMICS LEARNING WITH ACTION CONDITIONAL SSM


For the inverse dynamics case  we want to learn a model  where o⃗_t+1 is the desired next observation for time step t + 1 and â⃗̂_t is the predicted action, i.e., the one to be applied. We introduce an action decoder which decodes the latent posterior (z⃗_t^+, Σ⃗_t^+ ) and estimates the action required to move to the desired next observation. The action decoder also gets information regarding the next observation as an input.
During training this corresponds to the next observation in the data. 
For control a desired next observation is used to obtain the action required to reach that observation. The joint angles and velocities are treated as observations in our inverse dynamics experiments.



As seen in Figure <ref>, instead of merely learning the inverse dynamics, our approach learns inverse and forward dynamics simultaneously by feeding back the executed action to the action-conditional predict stage, which predicts forward in latent space. 



Note that the action decoder also gets the same action as the target during training.
However, the model sees the true action only after the prediction since the action decoder works with the posterior estimate (z⃗_t^+, Σ⃗_t^+ ).  


Hence, the prediction of the inverse dynamics model for the current time step is made independent of this feedback.
We found that enforcing this causal feedback, a necessary structural component of the Bayesian network of the underlying latent dynamical system, improves the performance of the inverse model, as seen in Figure <ref>.









  
Loss And Training.
The dual output architecture for our inverse model leads to two different loss functions for the action and observation decoders which are optimized jointly.
Our experiments showed that the loss of the forward model is an excellent auxiliary loss function for the inverse model and learning this implicit forward model jointly with the inverse model results in much better performance for the inverse model. 
We assume the reasons for this effect is that the forward model loss is providing crucial gradient information to form an informative latent state representation that is also useful for the inverse model.

In order to deal with high frequency data we again chose to predict the normalized differences to the previously executed action, i.e. â⃗̂_t = a⃗_t-1 + dec_action(z⃗_t^+). 
Here a⃗_t-1 is the executed action at t-1 and dec_action(z⃗_t^+) the output of the actual action decoder network. 
The combined loss function for the inverse dynamic case is given by 










 
    ℒ_inv = √(1T∑_i=1^T ∥(a⃗_t+1 - a⃗_t) - dec_action(z⃗_t^+) ∥^2)    + 
         λℒ_fwd

where 

λ chooses the trade-off between the inverse model loss and the forward model loss. The value of λ is chosen via hyperparameter optimization using GPyOpt gpyopt2016.
Note that for the inverse model,  the actions are decoded based on the posterior mean of the current time step while for the forward model the observations are decoded from the prior mean of the next time step. 


We evaluated the performance of the proposed method for inverse dynamics learning on two real robots, Franka Emika Panda and Barrett WAM. Barret WAM is a robot with direct cable drives. The direct cable drives produce high torques, generating fast and dexterous movements but yield complex dynamics. Rigid-body dynamics cannot model this complex dynamics accurately due to the variable stiffness and lengths of the cables nguyen2011incremental. 
   



Joint Torque Prediction Task. We benchmark the representational capability of the latent state posterior, z⃗_t^+, of ac-RKN in accurately modelling the inverse dynamics of these robots in comparison to deterministic models like LSTMs and FFNN. It is clear from <ref> and <ref> that ac-RKN learns highly accurate models of these in contrast with other data-driven methods. This highly precise modelling is often a requirement for high fidelity and compliant robotic control.   


























 
Impact of Action Feedback. We also perform an ablation study with and without the action feedback for the prediction step in the latent dynamics. As seen in the Figure <ref>, the action feedback always results in better representational capability as this helps the implicit forward model in making better predictions by taking into account its causal effect on the state transitions. 


 
Comparison to Analytical Models. Finally, we make a comparison with the analytical model of the Panda robot for inverse dynamics. Please refer to Appendix C for more details on the analytical model. As evident from Figure <ref> analytical models gave predictions with much lesser accuracy in comparison to ac-RKN, as it does not consider unmodelled effects such as joint and link flexibilities, backlash, stiction and actuator dynamics.  In such cases, the robot would continuously have to track its current position and compensate the errors with high-gain feedback control thus making it dangerous to interact with the real world and impossible to work in human-centred environments.




§ IMPLEMENTATION DETAILS




 §.§ Locally Linear Transition Model
 The state transitions in the predict stage of the Kalman filter is governed by a locally linear transition model. To obtain a locally linear transition model, the RKN learns K constant transition matrices A⃗^(k) and combines them using state dependent coefficients α^(k)(z⃗_⃗t⃗), i.e.,
A⃗_t = ∑_k=0^K α^(k)(z⃗_⃗t⃗) A⃗^(k).
A small neural network with softmax output is used to learn α^(k). Each A⃗^(k) is designed to consist of four band matrices as in <cit.> in order to reduce the number parameters without affecting the performance.




§ DETAILS OF RIGID BODY DYNAMICS MODEL

The analytical model for Franka Emika Panda is a rigid-body dynamics model that was identified in its so-called base parameters Khalil1991. Due to the friction compensation in the joints, we observed that the viscous friction is negligible, whereas the observed Coulomb friction is very small yet included in our parameterization. Which results is a model with 50 parameters. 
The base parameterization is computed based on provided kinematic properties of the robotic arm and provides a linear relation between the base parameters and joint torques for a given set of joint positions, velocities and accelerations. 

Due to this linearity, the regression problem can be solved using a linear least-squares method, although additional linear matrix inequality constraints must be fullfilled to ensure that the resulting parameters are physically realizable Sousa2019,reuss2022end.
In order to perform forward simulation of the robot dynamics, we numerically solve an initial value problem for the implicit set of differential equations defined by base parameterization of the rigid-body model.
Note that, the model does not parameterize actuator dynamics, nor does it model joint flexibilities, link flexibilities, or stiction. The focus here is to provide a reference baseline to show which effects the acRKN captures in comparison to a text-book robot model.


CHAPTER: APPENDIX: HIDDEN PARAMETER SSM



§ PROOF FOR GAUSSIAN IDENTITY <REF>


This section provides a proof for Gaussian Identity <ref>. First we derive an expression for the joint distribution p(u,v,y). 

[Joint Gaussian Distribution]gijointgauss2
If u∼𝒩(μ_u + b,Σ_u) and v ∼𝒩(μ_v,Σ_v) are normally distributed independent random variables and if conditional distribution p(y|u,v) = 𝒩(Au + b + Bv, Σ), the joint distribution has an expression as follows:
 
([ 𝐮; 𝐯; 𝐲 ])  ∼𝒩(([            μ_u;            μ_v; 𝐀μ_u+ b + 𝐁μ_v ]),([                 Σ_u                   0              Σ_u𝐀^⊤;                   0                 Σ_v              Σ_v𝐁^⊤;              𝐀Σ_u^⊤              𝐁Σ_v^⊤ 𝐀Σ_u𝐀^⊤+𝐁Σ_v𝐁^⊤ + Σ ]))


Proof for Identity <ref> 

Let displacement of a variable  𝐮 be denoted by Δ𝐮=𝐮-⟨𝐮⟩.

Since 𝐮 and 𝐯 are independent, covariances ⟨Δ𝐮Δ𝐯^⊤⟩=0.

We can write 𝐲=𝐀u + b + 𝐁v + ϵ, where ϵ∼𝒩(0,Σ) and b is a constant. Then we have covariance ⟨Δ𝐮Δ𝐲^⊤⟩=⟨Δ𝐮(𝐀Δ𝐮+𝐁Δ𝐯 + Δϵ)^⊤⟩= ⟨Δ𝐮Δ𝐮^⊤⟩𝐀^⊤+⟨Δ𝐮Δ𝐯^⊤⟩𝐁^⊤ +  ⟨Δ𝐮Δϵ^⊤⟩. Since ⟨Δ𝐮Δ𝐯^⊤⟩=⟨Δ𝐮Δϵ^⊤⟩=0 we therefore have ⟨Δ𝐮Δ𝐲^⊤⟩=Σ_u𝐀^⊤. 

The derivations for covariances ⟨Δ𝐯Δ𝐲^⊤⟩ follows similarly and the corresponding covariance has the expression, ⟨Δ𝐯Δ𝐲^⊤⟩=Σ_v𝐁^⊤.

Similarly, ⟨Δ𝐲Δ𝐲^⊤⟩= ⟨(𝐀Δ𝐮+𝐁Δ𝐯 + Δϵ)(𝐀Δ𝐮+𝐁Δ𝐯 + Δϵ)^⊤⟩=𝐀⟨Δ𝐮Δ𝐮^⊤⟩𝐀^⊤+𝐁⟨Δ𝐯Δ𝐯^⊤⟩𝐁^⊤ + ⟨ΔϵΔϵ^⊤⟩ = 𝐀Σ_u𝐀^⊤+𝐁Σ_v𝐁^⊤ +Σ. The result follows.

[Gaussian Marginalization]gigeneralgaussmarg

If

    ([ 𝐮; 𝐯; 𝐲 ]) ∼𝒩(([ μ_u; μ_v; μ_z ]),([   Σ_uu   Σ_uv   Σ_uy; Σ_uv^⊤   Σ_vv  Σ_v y; Σ_uy^⊤ Σ_vy^⊤   Σ_yy ]))

then marginal over y is given as p(y) = ∫_u,v p(y|u,v)p(u)p(v)dudv = 𝒩(μ_y, Σ_yy)


Proof For Identity <ref> 

We refer to <cit.> for the derivation, which requires calculation of the Schur complement as well as completing the square of the Gaussian p.d.f. to integrate out the variable. The given derivation bishop2006pattern for two variable multivariate Gaussians can be extended to 3 variable case WLOG.


Proof For Identity <ref> is immediate from Identity <ref> and Identity <ref>. 




§ ADDITIONAL EXPERIMENTS




 §.§ Comparison To Soft Switching Baseline


We also implement a soft-switching baseline similar to <cit.>, where a soft mixture of dynamics is implemented in the latent transition dynamics (Kalman time update).  Similar to <cit.>, we now globally learn K constant transition matrices A^(k) and control matrices B^(k). An interpolation is done between these using a “dynamics parameter network" fraccaro2017disentangled α_t= α_t(𝐰_0: t-1). The dynamics parameter network is implemented with a recurrent neural network with LSTM cells that takes at each time step the mean of the encoded observation w_t as input and recurses 𝐝_t= LSTM(𝐰_t-1, 𝐝_t-1) and α_t=softmax(𝐝_t). The output of the dynamics parameter network is weights that sum to one, ∑_k=1^Kα_t^(k)(𝐰_0: t-1)=1.
These weights choose and interpolate between K different operating modes:


    𝐀_t=∑_k=1^Kα_t^(k)(𝐰_0: t-1) 𝐀^(k),   𝐁_t=∑_k=1^Kα_t^(k)(𝐰_0: t-1) 𝐁^(k)


Authors interpret the weighted sum as a soft mixture of K different Linear Gaussian SSMs whose time-invariant matrices are combined using the time varying weights α_t. In practice, each of the K sets {𝐀^(k), 𝐁^(k)} models different/changing dynamics, that will dominate when the corresponding α_t^(k) is high. 

Figure <ref> compares HiP-RSSM with different recurrent architectures including the soft-switching baseline. As seen in figure <ref>, HiP-RSSM clearly outperforms the soft-switching baseline both in terms of convergence speed and also mult-step ahead predictions. More details regarding the multi-step ahead training procedure can be found in appendix <ref>.





 §.§ Multi-Step Ahead Predictions

 
In figure <ref>, we compare the results of multi-step ahead predictions for upto 50 steps of HiP-RSSM with recurrent baselines like RKNbecker2019recurrent and LSTM. Inorder to train the recurrent models for multi step ahead predictions, we removed three-quarters of the observations from the temporal sequence and tasked the models with imputing those missing observations, only based on the knowledge of available actions/control commands, i.e., we train the models to perform action conditional future predictions to impute missing observations. The imputation employs the model for multi-step ahead predictions in a convenient way shaj2020action. One could instead also go for a dedicated multi-step loss function as in approaches like <cit.>.

As seen in figure <ref>, HiP-RSSM clearly outperforms contemporary recurrent models for multi-step ahead prediction tasks since it takes into account additional causal factors of variation (slopes of the terrain for this robot) in the latent dynamics in an unsupervised manner.



§ IMPLEMENTATION DETAILS





 §.§ Context Set Encoder and Latent Task Representation
 The HiP-RSSM maps a set of previous interaction histories, {o_t,n^l,a_t,n^l,o_t+1,n^l}_n=1^N, to a set of latent features and an estimate of uncertainty in those features, {r_n^l,(σ_n^l)^2}_n=1^N, using a context encoder. We use a feed-forward neural network as the encoder in all of our experiements since we deal with high dimensional vectors as our observations. However, depending upon the nature of observations, we could use different encoder architectures.

The set of latent features and uncertainties, {r_n^l,(σ_n^l)^2}_n=1^N, are further aggregated in a probabilistically principled manner using bayesian aggregation operator discussed in <ref> to get a gaussian latent task variable, with a mean (μ_l) and diagonal covariance (σ_l). Intuitively the context encoder learns to weight the contribution from each observation in the context set based on bayesian priniciples and emits a probabilistic representation of the latent task.

 

 §.§ Latent Task Transformation Model

 
 To achieve latent task conditioning within the recurrent cell, we include  a task transformation model (c), in addition to the locally linear transition model A_t and control model b in the time update stage (section <ref>). 
Though in section <ref>, we used the notation for a linear task transformation matrix, C, to motivate the additive interaction of latent task variables, μ_l and σ_l in the latent space, the task transformation function can be designed in several ways, i.e.:


  (i) Linear: c =C, where C is a linear transformation matrix. The corresponding time update equations are given below:

    z_t^-    =A_t-1z_t-1^+ + b(a_t) + Cμ_l, 
    Σ_t^-    = A_t-1Σ_t-1^+A_t-1^T + C(I·σ_l)C^T + Σ_trans.



  (ii) Locally-Linear: c = C_t, where C_t = ∑_k=0^K β^(k)(z_t) C^(k) is a linear combination of k linear control models C^(k). A  small  neural  network  with softmax  output  is  used  to  learn β^(k). The corresponding time update equations are given below: 

    z_t^-    =A_t-1z_t-1^+ + b(a_t) + C_tμ_l, 
    Σ_t^-    = A_t-1Σ_t-1^+A_t-1^T + C_t(I·σ_l)C_t^T + Σ_trans.


  (iii) Non-Linear: c = f, where f(.) can be any non-linear function approximator. We use a multi-layer neural network regressor with ReLU activations, which transforms the latent task moments μ_l and σ_l directly into the latent space of the state space model via additive interactions. The corresponding time update equations are given below:

    z_t^-    =A_t-1z_t-1^+ + b(a_t) + f(μ_l), 
    Σ_t^-    = A_t-1Σ_t-1^+A_t-1^T + f(σ_l)+ Σ_trans.

 

In our ablation study (Figure <ref>), for the linear and locally linear task transformation models, we assume that the dimension of the latent context variable l and the latent state space z_t are equal. This allows us to work with square matrices which are more convenient. For the non-linear transformation we are free to choose the size of the latent context variable.  However for a fairer comparison we keep the dimension of latent task variable to be similar in all three cases. We choose the non-linear task transformation model in HiP-RSSM architecture as this gave the best performance in practice.


CHAPTER: APPENDIX: MULTI TIME SCALE SSM



§ PROOFS AND DERIVATIONS


[7]r.29
[b]0.28
   width=
            

Graphical Model For Bayesian conditioning with N observations.



In the following sections vectors are denoted by a lowercase letter in bold, such as "v", while Matrices as an uppercase letter in bold, such as "M". I denotes identity matrix and 0 represents a matrix filled with zeros. For any matrix M, m denotes the corresponding vector of diagonal entries. Also, ⊙ denotes the elementwise vector product and  ⊘ denotes an elementwise vector division.



 §.§ Proof For Bayesian Conditioning As Permutation Invariant Set Operations (Identity <ref>)



*

Note that Σ_l is the covariance matrix which is the inverse of the precision matrix Λ_l. Due to the observation model assumption H=[I,0], they take block diagonal form, 
    Σ_l = [[ Σ_l^u Σ^s_l; Σ^s_l Σ^l_l ]],  with Σ_u = diag(σ_l^u),  Σ_l = diag(σ_l^l)  and Σ_s = diag(σ_l^s).
  



  
Proof:



  
Case 1 (Single Observation): Before deriving the update rule for N conditionally iid observations, let us start with a simpler case consisting of a single observation r. If the marginal Gaussian distribution for the latent variable l takes the form p(𝐥) =𝒩(𝐥|μ, Λ^-1) and the conditional Gaussian distribution for the single observation r given l has the form
, p(𝐫|𝐥) =𝒩(𝐫|𝐇𝐥+𝐛, 𝐋^-1). Then the posterior distribution over 𝐥 can be obtained in closed form as, 

    p(𝐥|𝐫)     =𝒩(𝐥|Σ{𝐇^T𝐋(𝐫-𝐛)+Λμ}, Λ^-1) ,where Λ=(Λ+𝐇^T𝐋𝐇).

We refer to Section 2.3.3 of <cit.>, to the proof for this standard result.



  
Case 2 (Set Of Observations): Now instead of a single observation, we wish to derive a closed form solution for the posterior over latent variable l∈ℝ^2d, given a set of N conditionally i.i.d observations r = {r_i}_i=1^N. Here each element r_i∈ℝ^d of the set r is assumed to to have an observation model H=[I,0]. In the derivation, we represent the set of N observations as a random vector 
    r = [ [ r_1; r_2;   .;   .; r_N ]]_Nd × 1.
 
Since each observation in the set r are conditionally independent, we denote the conditional distribution over the context set as r|𝐥∼𝒩(H̅l,  Σ_r), where the diagonal covariance matrix has the following form: 
    Σ_r= [ [ diag(σ_r_1),           0,           0,          ..,            0;           0, diag(σ_r_2),           0,          ..,            0;           .,           .,           .,          ..,            .;          . ,           .,           .,          ..,            .;           0,           0,           0,          ..,  diag(σ_r_N) ]]_Nd × Nd.

The corresponding observation model  H̅ is

    H̅ =  [ [ H; H; .; .; H ]]_Nd × 2d = [ [ I , 0; I , 0;  ., .; . , .; I , 0 ]]_Nd × 2d.


Now given the prior over the latent task variable l∼𝒩(μ_l^-, Σ_l^-), the parameters of the posterior distribution over the task variable, p(l|r) ∼𝒩(μ_l^+, Λ_l^+), can be obtained in closed-form substituting in Equation (<ref>) as follows.

    Λ_l^+     =  (Σ_l^+)^-1
       = Σ_l^-1+ H̅ ^T Σ_rH̅
        = Σ_l^-1+ [ [ diag(σ_r_1),  diag(σ_r_2),  diag(σ_r_3), . , . ,   diag(σ_r_N);                                 0 ,     0,     0,   . , . ,  0;                                                                ]]_2d× ndH̅
        =λ_l^- + [ [ diag(∑_i=1^n1/σ_r_i),                     0;                   0 ,                     0;                       ]]_2d× 2d
    μ_l^+     = μ_l^- + (Λ^+)^-1H̅ ^T ( σ_r^-2I)(y-H̅μ_x) 
        = μ_l^-+ Σ^+ H̅( σ_r^-2I) (y- H̅μ_x) 
        = μ_l^-+ Σ^+[ [ σ_r_1^-2I , σ_r_2^-2I, σ_r_3^-2I , . , . ,  σ_r_n^-2I;                        0 ,     0,     0,   . , . ,  0;                                                       ]] (y-H̅μ_x) 
        =μ_l^- + [ [  σ_l^u+,   σ_l^s+; σ_l^s+ ,   σ_l^l+;          ]]  [ [ ∑_n=1^N(𝐫_𝐧-μ^u,-_l) ⊙1/σ_i;                           0;                             ]]  
        =μ_l^- + [ [ σ_l^u+; σ_l^s+;        ]]  ⊙[ [ ∑_i=1^N(𝐫_𝐢-μ^u,-_l) ⊙1/σ_r_i; ∑_i=1^N(𝐫_𝐧-μ^u,-_l) ⊙1/σ_r_i;                               ]]


Here μ_l^+ is the posterior mean and Λ_l^+ is the posterior precision matrix.




 §.§ Derivation For Matrix Inversions as Scalar Operations



Consider a block  matrix of the following form A = [[ diag( a^u) diag( a^s); diag( a^s) diag( a^l) ]]. Then inverse A^-1 =  B can be calculated using scalar operations and is given as, B = [[ diag( b^u) diag( b^s); diag( b^s) diag( b^l) ]] where,

    b^u     = a_l ⊘ (  a_u ⊙ a_l-  a_s ⊙ a_s )
     b^s     = -  a_s ⊘ (  a_u ⊙ a_l-  a_s ⊙ a_s )
     b^l     = a_u ⊘ (  a_u ⊙ a_l-  a_s ⊙ a_s )
.





  
Proof: To prove this we will use the following matrix identity of a partitioned matrix from <cit.>, which states

    ([ A B; C D ])^-1=([                M          -MBD^-1;          -D^-1CM D^-1+D^-1CMBD^-1 ])

where M is defined as

    M=(A-BD^-1C)^-1.
 Here M is called the Schur complement of the Matrix on the left side of Equation <ref>. The algebraic manipulations to arrive at scalar operations in Equation <ref> are straightforward.



 §.§ Proof for Permutation Invariance (Theorem <ref>)
 


*

The posterior distribution p(θ| X), defined by Bayes' theorem as p(θ| X) = p(X |θ)p(θ)/p(X), relies on the likelihood p(X |θ) = ∏_i=1^N p(x_i |θ). Under any permutation π, the permuted set X^π = {π(x_1), …, π(x_N)} preserves the likelihood since p(X^π|θ) = ∏_i=1^N p(π(x_i) |θ) = ∏_i=1^N p(x_i |θ) = p(X |θ). The prior p(θ) and the marginal likelihood p(X) are invariant under permutation. Therefore, p(θ| X^π) = p(θ| X), establishing the permutation invariance of the posterior.




 §.§ Proof for Gaussian Marginalization (Identity <ref>)
 

*
  
The proof is a staright forward extension to the Identity <ref> derived in Chapter <ref>.  First, we derive an expression for the joint distribution p(u_1,u_2,..,u_N,y).

[Joint Distribution]gigaussjoint
 If u_1∼𝒩(μ_1,Σ_1) and u_2∼𝒩(μ_2,Σ_2) are normally distributed independent random variables and if conditional distribution p(y|u_1,u_2) = 𝒩(A_1u_1 + A_2u_2, Σ), the joint distribution has an expression as follows:

    ([ 𝐮_1; 𝐮_2;   ⋮; 𝐮_𝐍;   𝐲 ]) ∼𝒩(
    ([             μ_1;             μ_2;               ⋮;             μ_N; ∑_n=1^N 𝐀_n μ_n ]),
    ([                       Σ_1                         0                         ⋯                  𝐀_1^⊤Σ_1;                         0                       Σ_2                         ⋯                  𝐀_2^⊤Σ_2;                         ⋮                         ⋮                         ⋱                         ⋮;                   Σ_1 𝐀_1                   Σ_2 𝐀_2                         ⋯ ∑_n=1^N 𝐀_n Σ_n 𝐀_n^⊤ + Σ ])
    )




We can write 𝐲= ∑_n=1^N A_nu_n + ϵ, where ϵ∼𝒩(0,Σ) and b is a constant. 

Let displacement of a variable  𝐮 be denoted by Δ𝐮=𝐮-⟨𝐮⟩.

Since 𝐮_𝐢 and 𝐮_𝐣 are independent ∀ i ≠ j, the covariances

    ⟨Δ𝐮_𝐢Δ𝐮_𝐣^⊤⟩=0, ∀ i ≠ j .


Similarly, 
    ⟨Δ𝐮_𝐢Δϵ^⊤⟩=0, ∀ i.


For any i, we have the covariance

    ⟨Δ𝐮_i Δ𝐲^⊤⟩    = ⟨Δ𝐮_i (∑_n=1^N 𝐀_n Δ𝐮_n + Δϵ)^⊤⟩
        = ⟨Δ𝐮_i Δ𝐮_i^⊤⟩𝐀_i^⊤ + ∑_j=1, j ≠ i^N ⟨Δ𝐮_i Δ𝐮_j^⊤⟩𝐀_j^⊤ + ⟨Δ𝐮_i Δϵ^⊤⟩.



Using equations <ref> and <ref>, we therefore derive an expression for the corresponding covariance as:

    ⟨Δ𝐮_i Δ𝐲^⊤⟩ = Σ_i 𝐀_i^⊤.


Similarly, 

    ⟨Δ𝐲Δ𝐲^⊤⟩   = ⟨(∑_n=1^N 𝐀_n Δ𝐮_n + Δϵ) (∑_n=1^N 𝐀_n Δ𝐮_n + Δϵ)^⊤⟩
       = ∑_n=1^N 𝐀_n ⟨Δ𝐮_n Δ𝐮_n^⊤⟩𝐀_n^⊤ + ⟨ΔϵΔϵ^⊤⟩
       = ∑_n=1^N 𝐀_n Σ_n 𝐀_n^⊤ + Σ_ϵ.

The result follows.
[Marginalization]gigeneralgaussmarg

If

    ([ u_1; u_2;   ⋮; u_N;   y ]) ∼𝒩(([ μ_1; μ_2;   ⋮; μ_N; μ_y ]),([ Σ_11 Σ_12    ⋯ Σ_1N Σ_1y; Σ_21 Σ_22    ⋯ Σ_2N Σ_2y;    ⋮    ⋮    ⋱    ⋮    ⋮; Σ_N1 Σ_N2    ⋯ Σ_NN Σ_Ny; Σ_y1 Σ_y2    ⋯ Σ_yN Σ_yy ]))

then marginal over y is given as p(y) = ∫ p(y, u_1, u_2, . ., u_N)∏_i=1^N u_i = 𝒩(μ_y, Σ_yy)



Proof For Identity <ref> 
We refer to <cit.> for the derivation, which requires calculation of the Schur complement as well as completing the square of the Gaussian p.d.f. to integrate out the variable. The given derivation bishop2006pattern for two variable multivariate Gaussians can be extended to the case of N variables WLOG.


Proof For Identity <ref> is immediate from Identity <ref> and Identity <ref>.



§ IMPLEMENTATION DETAILS





 §.§ Inference In Slow Time Scale SSM





  §.§.§ Inferring Action Abstraction (sts-SSM)

Given a set of encoded primitive actions and their corresponding variances {α_k,t, ρ_k,t}_t=1^H, using the prior and observation model assumptions in Section 3.1.2 of main paper, we infer the latent abstract action p(α_k|α_k,1:H) = 𝒩(μ_α_k, Σ_α_k)  = 𝒩(μ_α_k, diag(σ_α_k)) as a Bayesian aggregation <cit.> of these using the following closed-form equations:
[11]r.31
[b]0.31
   width=
         
     
Generative model for the abstract action α_k. The hollow arrows are deterministic transformations leading to implicit distribution α_k,t using an action set encoder.




    σ_α_k    = ( ( σ_0)^⊖ + ∑_n=1^N ((ρ_k,t)^⊖) )^⊖, 
    μ_α_k    = μ_0 + σ_α_k⊙∑_n=1^N (α_k,t - μ_0) ⊘ρ_k,t


Here, ⊖, ⊙ and ⊘ denote element-wise inversion, product, and division, respectively. The update equation is coded as the “abstract action inference” neural network layer as shown in Figure <ref>. 



  §.§.§ Task Prediction (sts-SSM)



The goal of this step is to update the prior marginal over the latent task variable l_k, p(l_k|β_1:k-1,α_1:k), given the posterior beliefs from the time window k-1 and abstract action α_k.

Using the linear dynamics model assumptions from Section 3.3, we can use the following closed-form update equations to compute,
p(l_k|β_1:k-1,α_1:k) = 𝒩(μ_l_k^-,Σ_l_k^-), where

    [c]
    μ_l_k^-    =Xμ_l_k-1^+ + Yα_k
    Σ_l_k^-    = XΣ_l_k-1^+X^T + YΣ_α_kY^T +  S.


These closed-form equations are coded as the “task predict” neural net layer as shown in Figure <ref>.



  §.§.§ Task Update (sts-SSM)


In this stage, we update the prior over l_k using an abstract observation set {β_k,t}_t=1^H, to obtain the latent task the posterior 𝒩(μ_z_k,t^+,Σ_z_k,t^+) = 𝒩([[ μ_t^u+; μ_t^l+ ]],[[ Σ_t^u Σ^s_t; Σ^s_t Σ^l_t ]]^+), with Σ_l_k^u = diag(σ_l_k^u),  Σ_l_k^l = diag(σ_l_k^l)  and  Σ_l_k^s = diag(σ_l_k^s).

To do so we first invert the prior covariance matrix [[ Σ_l_k^u Σ_l_k^s; Σ_l_k^s Σ_l_k^l ]]^+ to the precision matrix [[ λ_l_k^u λ_l_k^s; λ_l_k^s λ_l_k^l ]]^+ for permutation invariant parallel processing. The posterior precision is then computed using scalar operations are follows, where only λ_l_k^u is changed by 

    λ_l_k^u+ = λ_l_k^u- + ∑_t=1^H  1 ⊘ν_k,t

 while λ_l_k^l+ = λ_l_k^l- and λ_l_k^s+ = λ_l_k^s- remain constant. The operator ⊘ denotes the element-wise division. The posterior precision is inverted back to the posterior covariance vectors σ_l_k^u+, σ_l_k^l+ and σ_l_k^s+. Now, the posterior mean μ_l,k^+ can be obtained from the prior mean μ_l,k^- as
 
    [c]
       μ_l,k^+=μ_l,k^- + [ [ σ_l_k^u+; σ_l_k^s+;          ]]  ⊙[ [ ∑_t=1^H(β_k,t-μ^u,-_l_k) ⊘ν_k,t; ∑_t=1^H(β_k,t-μ^u,-_l_k) ⊘ν_k,t;                                 ]].



The inversion between the covariance matrix and precision matrix can be done via scalar operations leveraging block diagonal structure as derived in Appendix <ref>. Figure <ref> shows the schematic of the task update layer.


 §.§ Inference In Fast Time Scale SSM


The inference in fts-SSM for a time-window k involves two stages as illustrated in Figure <ref>, calculating the prior and posterior over the latent state variable z_t. To keep the notation uncluttered, we will also omit the time-window index k whenever the context is clear as in Section 3.2.



  §.§.§ Task Conditional State Prediction (fts-SSM)


Following the assumptions of a task conditional linear dynamics as in Section 3.2 of the main paper, we obtain the prior marginal for p(z_k,t|w^k_1:t-1,a^k_1:t-1,β_1:k-1,α_1:k-1) = 𝒩(μ_z_k,t^-,Σ_z_k,t^-) in closed form, where

    μ_z_k,t^-    =Aμ_z_k,t-1^- + Ba_k,t-1 + Cμ_l_k^-, 
    Σ_k,t^-    = AΣ_k,t-1^+A^T + CΣ_l_k^-C^T + Q.




  §.§.§ Observation Update (fts-SSM)

In this stage, we compute the posterior belief p(z_k,t|w^k_1:t,a^k_1:t,β_1:k,α_1:k-1) = 𝒩(μ_z_k,t^-,Σ_z_k,t^-).  using the same closed-form update as in  <cit.>. The choice of the special observation model splits the state into two parts, an upper z_t^u and a lower part z_t^l, resulting in the posterior belief 𝒩(μ_z_k,t^-,Σ_z_k,t^-) = 𝒩([[ μ_t^u+; μ_t^l+ ]], [[ Σ_t^u Σ^s_t; Σ^s_t Σ^l_t ]]^+), with Σ_t^u = diag(σ_t^s),  Σ_t^l = diag(σ_t^l) and Σ_t^s = diag(σ_t^s). Thus, the factorization allows for only the diagonal and one off-diagonal vector of the covariance to be computed and simplifies the calculation of the mean and posterior to simple scalar operations. 

The closed-form equations for the mean can be expressed as the following scalar equations,

    z_t^+ = z_t^- +
    [[ σ^u,-_t; σ^l,-_t ]]
    ⊙[[ w_t - z^u,-_t; w_t - z^u,-_t ]]
    ⊘[[ σ_t^u,- + σ_t^obs; σ_t^u,- + σ_t^obs ]],


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

    σ^u,+_t    = σ^u,-_t ⊙σ^u,-_t ⊘( σ_t^u,- + σ_t^obs),  
    σ^s,+_t    = σ^u,-_t ⊙σ^s,-_t ⊘( σ_t^u,- + σ_t^obs), 
    σ^l,+_t    = σ^l, -_t - σ^s,-_t ⊙σ^s,-_t ⊘( σ_t^u,- + σ_t^obs),
,
where ⊙ denotes the elementwise vector product and  ⊘ denotes an elementwise vector division. 



 §.§ Modelling Assumptions




  §.§.§ Control Model
 To achieve action conditioning within the recurrent cell of fts-SMM, we include a control model b(a_k,t) in addition to the linear transition model A_t. b(a_k,t) =f(a_k,t), where f(.) can be any non-linear function approximator.  We use a multi-layer neural network regressor with ReLU activations <cit.>. 

However, unlike the fts-SSM where actions are assumed to be known and subjected to no noise, in the sts-SSM, the abstract action is an inferred latent variable with an associated uncertainty estimate. Hence we use a linear control model Y, for principled uncertainty propagation. 



  §.§.§ Transition Noise

We assume the covariance of the transition noise Q and S in both timescales to be diagonal. The noise is learned and is independent of the latent state.



§ METRICS USED FOR MEASURING LONG HORIZON PREDICTIONS




 §.§ Sliding Window RMSE

The sliding window RMSE (Root Mean Squared Error) metric is computed for a predicted trajectory in comparison to its ground truth. At each time step, the RMSE for each trajectory is determined by taking the root mean square of the differences between the ground truth and predicted values within a sliding window that terminates at the current time step. This sliding window, with a specified size, provides a smoothed localized assessment of prediction accuracy over the entire prediction length.
Mathematically, the sliding window RMSE at time step t is given by:

    RMSE(t)=√(1/W∑_i=t-W+1^t(gt_i-pred_i)^2)

where t is the current time step, W is the window size, and gt_i and pred_i are the ground truth and predicted values at time step i, respectively. The extension to multiple trajectories is straightforward and omitted to keep the notation uncluttered.



 §.§ Sliding Window NLL

The sliding window NLL (Negative Log-Likelihood) metric is computed for a predicted probability distribution against the true distribution. At each time step, the NLL is determined by summing the negative log-likelihood values within a sliding window that terminates at the current time step. This sliding window, with a specified size, provides a smoothed localized evaluation of prediction accuracy across the entire sequence.

Mathematically, the sliding window NLL at time step t is given by:

    NLL(t)=-1/W∑_i=t-W+1^t log𝒩(gt_i |predMean_i, predVar_i)

where t is the current time step, W is the window size. predMean_i, predVar_i, and gt_i represent the predicted mean, predicted variance, and the ground truth at time step i.



§ VISUALIZATION OF PREDICTIONS GIVEN BY DIFFERENT MODELS.


 In this section, we plot the multistep ahead predictions (mean and variance) by different models on 3 datasets on normalized test trajectories. Not that we omit NaN values in predictions while plotting.
 

 

  §.§.§ Franka Kitchen

  


  §.§.§ Hydraulic Excavator

 





  §.§.§ Mobile Robot

 



CHAPTER: APPENDIX: ROBOTS AND DATASET DETAILS



§ ROBOTS AND DATA USED FOR IN CHAPTER <REF> ON AC-SSM



The experiments are performed on data from four different robots. The details of robots, data, and data preprocessing are explained below: 


 §.§ Hydraulic Brokk 40 Robot Arm


Observation and Data Set: The data was obtained from a HYDROLEK–7W 6 degree–of–freedom manipulator with a continuous (360 degree) jaw rotation mechanism. We actuate the joints via hydraulic pistons, which are powered via an auxiliary output from the hydraulic pump. Thus learning the forward model is difficult due to inherent hysteresis associated with hydraulic control. For this robot, only one joint is moved at a time, so we have independent time series per joint. The joint data consists of measured joint positions and the input current to the controller of the joint sampled at 100Hz. 


Training Procedure: During training, we work with sequences of length 500. For the first 300 time steps those sequences consist of the full observation, i.e., the joint position and current. We give only the current signals in the remaining 200 time steps. The models have to impute the missing joint positions in an uninformed fashion, i.e., we only indicate the absence of a position by unrealistically high values. 





 §.§ Musculoskeletal Robot Arm



Observation and Data Set:  For this soft robot we have 4 dimensional observation inputs(joint angles) and 8 dimensional action inputs(pressures). We collected the data of a four DoF robot actuated by Pneumatic Artificial Muscles
(PAMs). The robot arm has eight PAMs in total with each DoF actuated by an antagonistic pair. The robot arm reaches high joint angle accelerations of up to 28, 000deg/s2 while avoiding dangerous joint limits thanks to the antagonistic actuation and limits on the air pressure ranges. The data
consists of trajectories collected while training with a model-free reinforcement learning algorithm to hit balls while playing table tennis. We sampled the data at 100Hz. The hysteresis associated with the pneumatic actuators used in this robot is challenging to model and is relevant to the soft robotics in general.


Training Procedure: During training, we randomly removed three-quarters of the states from the sequences and tasked the models with imputing those missing states, only based on the knowledge of available actions/control commands, i.e., we train the models to perform action conditional future predictions to impute missing states. The imputation employs the model for multi-step ahead predictions in a convenient way. One could instead go for a dedicated loss function as in approaches like <cit.>, <cit.> for long term predictions.



  §.§.§ Franka Emika Panda Robot Arm


Observation and Data Set: We collected the data from a 7 DoF Franka Emika Panda manipulator during free motion. We chose this task since the robot exhibits different dynamics behaviour due to electric actuators and high frequencies(1kHz). The raw joint positions, velocities and torques were recorded using Franka Interfaces while the joint accelerations were computed by finite differences on filtered velocity data (obtained using a zero-phase 8th-order digital Butterworth filter with a cut-off frequency of 5Hz). The observations for the forward model consist of the seven joint angles in radians, and the corresponding actions were joint Torques in Nm. While the inverse model use both joint angles and velocities as observations. The data was divided into train and test sets in the ratio 4:1. We divide the data into sequences of length 300 while training the recurrent models for forward dynamics and use sequences of length 50 for inverse dynamics.   
  

Training Procedure Forward Dynamics: Similar to the multi-step ahead training procedure in <ref>, during training we randomly removed three-quarters of the observations(joint angles) from the sequences and tasked the models with imputing those missing observations, only based on the knowledge of available actions/control commands.

Training Procedure Inverse Dynamics: The recurrent models (LSTM, ac-RKN) uses a similar architecture, as shown in Figure 3 of the main paper, except for the recurrent module. The hyperparameters including learning rate, latent state and observation dimensions, learning rate, control model architecture, action decoder architecture and regularization parameter for the joint forward-inverse dynamics loss function are searched via GpyOpt<cit.> and is mentioned in Appendix D. The observation encoder and decoder architecture is chosen to be of the same size across the models being compared. For all models, we use the joint positions and velocities as the observation input and differences to the next state as desired observation. The FFNN gets the current observation and desired observation as input and is tasked to predict the joint Torques directly(unlike differences in recurrent models) as in previous regression approaches<cit.>. 




  §.§.§ Barrett WAM Robot Arm


Observation and Data Set: The Barett task is based on a publicly available dataset comprising joint positions, velocities, acceleration and torques of a seven degrees-of-freedom real Barett WAM robot. The original training dataset (12, 000 data points) is split into sequences of length 98. Twenty-four out of the total 119 episodes are utilized for testing, whereas the other 95 are used for training. The direct cable drives which drive this robot produce high torques, generating fast and dexterous movements but yield complex dynamics.  Rigid-body dynamics cannot be model this complex dynamics due to the variable stiffness and lengths of the cables. 
  

Training Procedure Inverse Dynamics: The training procedure is repeated as in <ref>



§ ROBOTS AND DATASETS USED IN CHAPTER <REF> ON HIP-SSM




 §.§ Franka Emika Panda Robot Arm


Observation and Data Set: We collected the data from a 7 DoF Franka Emika Panda manipulator during free motion and while manipulating loads with weights 0kg (free motion), 0.5 kg, 1 kg, 1.5 kg, 2 kg and 2.5 kg. Data is sampled at high frequencies (1kHz). The training trajectories were motions with loads 0kg(free motion), 1kg, 1.5kg, 2.5 kgs, while the testing trajectories contained motions with loads of 0.5kg and 2 kgs. The observations for the forward model consist of the seven joint angles in radians, and the corresponding actions were joint Torques in Nm. We divide the data into sequences of length 600 while training the recurrent models for forward dynamics, with 300 time-steps (corresponding to 300 milli-seconds) used as context set and rest 300 is used for the recurrent time-series modelling.   
  

Training Procedure: For the fully observable case, we trained HiP-RSSM for one-step ahead prediction using an RMSE loss. Similar to the training procedure for partially observable case as in <ref>, during training we randomly removed half of the observations(joint angles) from the sequences and tasked the models with imputing those missing observations, only based on the knowledge of available actions/control commands.

Other recurrent baselines (RKN, LSTM, GRU) are trained in a similar fashion except that, we dont maintain a context set of interaction histories during training/inference.



 §.§ Wheeled Mobile Robot


Observation and Data Set: We collected 50 random trajectories from a Pybullet simulator a wheeled mobile robot traversing terrain with sinusoidal slopes. Data is sampled at high frequencies (500Hz). 40 out of the 50 trajectories were used for training and the rest 10 for testing. The observations consists of parameters which completely describe its location and orientation of the robot. The observation of the robot at any time instance t consists of the following features:.
    [ o_t=[x, y, z, cos (α), sin (α), cos (β);              sin (β), cos (γ), sin (γ)] ]

where,
x, y, z - denote the global position of the Center of Mass of he robot,
α, β, γ- Roll, pitch and yaw angles of the robot respectively, in the global frame of reference sonker2020adding.  We divide the data into sequences of length 300 while training the recurrent models for forward dynamics, with 150 time-steps (corresponding to 300 milli-seconds) used as context set and rest 150 is used for the recurrent time-series modelling.   
  

Training Procedure: For the fully observable case, we trained HiP-RSSM for one-step ahead prediction using an RMSE loss. Similar to the training procedure for partially observable case as in <ref>, during training we randomly removed half of the observations from the sequences and tasked the models with imputing those missing observations, only based on the knowledge of available actions/control commands.

Other recurrent baselines (RKN, LSTM, GRU) are trained in a similar fashion except that, we dont maintain a context set of interaction histories during training/inference.



§ ROBOTS AND DATA USED IN CHAPTER <REF> ON MULTI TIME SCALE SSM


In all datasets, we only use information about agent/object positions and we mask out velocities to create a partially observable setting. All datasets are subjected to a mean zero, unit variance normalization during training. During testing, they are denormalized after predictions. The details of the different datasets used are explained below:



 §.§ D4RL Datasets


Details: We use a set of 3 different environments/agents from D4RL dataset <cit.>, which includes the HalfCheetah, Franka Kitchen and Maze2D (medium) environment. (a) HalfCheetah: We used 1000 suboptimal trajectories collected from a policy trained to approximately 1/3 the performance of the expert. The observation space consists of 8 joint positions and the action space consists of 6 joint torques collected at 50 Hz frequency. 800 trajectories were used for training and 200 for testing. For the long horizon task, we used 1.2 seconds (60 timesteps) as context and tasked the model to predict 6 seconds (300 timesteps) into the future. (b) Franka Kitchen: The goal of the Franka Kitchen environment is to interact with the various objects to reach a desired state configuration. The objects you can interact with include the position of the kettle, flipping the light switch, opening and closing the microwave and cabinet doors, or sliding the other cabinet door. We used the "complete" version of the dataset and collected 1000 trajectories where all four tasks are performed in order. The observation space consists of 30 dimensions (9 joint positions of the robot and 21 object positions). The action space consists of 9 joint velocities clipped between -1 and 1 rad/s. The data was collected at a 50 Hz frequency. 800 trajectories were used for training and 200 for testing. For the long horizon task, we used 0.6 seconds (30 timesteps) as context and tasked the model to predict 2.7 seconds (135 timesteps) into the future. The dataset is complex due to multi-task, multi-object interactions in a single trajectory.
(c) Medium Maze: We used 20000 trajectories from a 2D Maze environment, where each trajectory consists of a force-actuated ball (along the X and Y axis) moving to a fixed target location. The observation consists of as the (x, y) locations and a 2D action space. The data is collected at 100 Hz frequency. 16000 trajectories were used for training and 4000 for testing. For the long horizon task, we used 0.6 seconds (60 timesteps) as context and tasked the model to predict 3.9 seconds (390 timesteps) into the future. Rendering of the three environments is shown in Figure <ref>.







 §.§ Hydraulic Excavator



Details: We collected the data from a wheeled excavator JCB Hydradig 110W show in Figure <ref>. The data was collected by actuating the boom and arm of the excavator using Multisine and Amplitude-Modulated Pseudo-Random Binary Sequence
(APRBS) joystick signals with safety mechanisms in place. A total of 150 mins of data was collected at a frequency of 100 Hz. of which was used as a training dataset and the rest as testing. The observation space consists of the boom and arm positions, while the joystick
signals are chosen as actions. For the long horizon task we used 1.5 seconds (150 timesteps) as context and tasked the model to predict 12 seconds (1200 timesteps) into the future.  
  






 §.§ Panda Robot With Varying Payloads


Details: We collected the data from a 7 DoF Franka Emika Panda manipulator during free motion and while manipulating loads with weights 0kg (free motion), 0.5 kg, 1 kg, 1.5 kg, 2 kg and 2.5 kg. The robot used is shown in Figure <ref>. Data is sampled at a frequency of 100 Hz. The training trajectories were motions with loads of 0kg(free motion), 1kg, 1.5kg, and 2.5 kgs, while the testing trajectories contained motions with loads of 0.5 kg and 2 kg. The observations for the forward model consist of the seven joint angles in radians, and the corresponding actions were joint Torques in Nm. For the long horizon task we used 0.6 seconds (60 timesteps) as context and tasked the model to predict 1.8 seconds (180 timesteps) into the future.  
  






 §.§ Wheeled Mobile Robot


Observation and Data Set: We collected 50 random trajectories from a Pybullet simulator a wheeled mobile robot traversing terrain with slopes generated by a mix of sine waves as opposed to the sine wave terrain for experiments used in Chapter <ref>, making this more challenging. Data is sampled at high frequencies (500Hz). 40 out of the 50 trajectories were used for training and the rest 10 for testing. The observations consist of parameters which completely describe the location and orientation of the robot. The observation of the robot at any time instance t consists of the following features:

    [ o_t=[x, y, z, cos (α), sin (α), cos (β);              sin (β), cos (γ), sin (γ)] ]

where,
x, y, z - denote the global position of the Center of Mass of the robot,
α, β, γ- Roll, pitch and yaw angles of the robot respectively, in the global frame of reference sonker2020adding.  For the long horizon task we used 0.6 seconds (150 timesteps) as context and tasked the model to predict 3 seconds (750 timesteps) into the future.   
  


CHAPTER: APPENDIX: HYPERPARAMETERS



§ HYPERPARAMETERS: CHAPTER <REF>




 §.§ Pneumatic Musculoskeltal Robot Arm




Encoder (ac-RKN,RKN,LSTM): 1 fully connected + linear output  (elu + 1) 

	
  * Fully Connected 1: 120, ReLU
 
Observation Decoder (ac-RKN,RKN,LSTM): 1 fully connected + linear output:

	
  * Fully Connected 1: 120, ReLU
 
 
Transition Model (ac-RKN,RKN):
bandwidth: 3, number of basis: 15

    
  * α(z⃗_t): No hidden layers - softmax output

Control Model (ac-RKN):
3 fully connected + linear output

	
  * Fully Connected 1: 120, ReLU
	
  * Fully Connected 2: 120, ReLU
	
  * Fully Connected 3: 120, ReLU
 

Architecture For FFNN Baseline
2 fully connected + linear output

	
  * Fully Connected 1: 6000, ReLU
	
  * Fully Connected 2: 3000, ReLU


Dropout Regularization - 0.512

Learning Rate - 1.39e-2

Optimizer Used: Adam Optimizer



 §.§ Hydraulic Brokk-40 Robot Arm




Encoder (ac-RKN,RKN,LSTM,GRU): 1 fully connected + linear output  (elu + 1) 

	
  * Fully Connected 1: 30, ReLU
 
Observation Decoder (ac-RKN,RKN,LSTM,GRU): 1 fully connected + linear output:

	
  * Fully Connected 1: 30, ReLU
 
 
Transition Model (ac-RKN,RKN):
bandwidth: 3, number of basis: 32

    
  * α(z⃗_t): No hidden layers - softmax output

Control Model (ac-RKN):
1 fully connected + linear output

	
  * Fully Connected 1: 120, ReLU
 



 §.§ Franka Emika Panda - Forward Dynamics Learning


Encoder (ac-RKN,RKN,LSTM,GRU): 1 fully connected + linear output  (elu + 1) 

	
  * Fully Connected 1: 120, ReLU
 
Observation Decoder (ac-RKN,RKN,LSTM,GRU): 1 fully connected + linear output:

	
  * Fully Connected 1: 240, ReLU
 
 
Transition Model (ac-RKN,RKN):
bandwidth: 3, number of basis: 15

    
  * α(z⃗_t): No hidden layers - softmax output

Control Model (ac-RKN):
3 fully connected + linear output

	
  * Fully Connected 1: 30, ReLU
	
  * Fully Connected 2: 30, ReLU
	
  * Fully Connected 3: 30, ReLU
 

Architecture For FFNN Baseline - Forward Dynamics
3 fully connected + linear output

	
  * Fully Connected 1: 1000, ReLU
	
  * Fully Connected 2: 1000, ReLU
	
  * Fully Connected 3: 1000, ReLU


Dropout Regularization - 0.1147

Learning Rate - 8.39e-3

Optimizer Used: SGD Optimizer



 §.§ Franka Emika Panda - Inverse Dynamics Learning



Encoder (ac-RKN,RKN,LSTM): 1 fully connected + linear output (elu + 1) 

    
  * Fully Connected 1: 120, ReLU
 
Observation Decoder (ac-RKN,RKN,LSTM): 1 fully connected + linear output:

    
  * Fully Connected 1: 240, ReLU
 

Action Decoder (ac-RKN,RKN,LSTM): 1 fully connected + linear output:

    
  * Fully Connected 1: 512, ReLU
 
 
Transition Model (ac-RKN,RKN):
bandwidth: 3, number of basis: 15

    
  * α(z⃗_t): No hidden layers - softmax output  


Control Model (ac-RKN):
1 fully connected + linear output

    
  * Fully Connected 1: 45, ReLU
 

Architecture For FFNN Baseline - Inverse Dynamics
3 fully connected + linear output

    
  * Fully Connected 1: 500, ReLU
    
  * Fully Connected 2: 500, ReLU
    
  * Fully Connected 3: 500, ReLU


Dropout Regularization - 0.563 

Learning Rate - 1.39e-2 

Optimizer Used: SGD Optimizer



 §.§ Barrett WAM - Inverse Dynamics Learning



Encoder (ac-RKN,RKN,LSTM): 1 fully connected + linear output (elu + 1) 

    
  * Fully Connected 1: 120, ReLU
 
Observation Decoder (ac-RKN,RKN,LSTM): 1 fully connected + linear output:

    
  * Fully Connected 1: 240, ReLU
 

Action Decoder (ac-RKN): 2 fully connected + linear output:

    
  * Fully Connected 1: 256, ReLU
    
  * Fully Connected 1: 256, ReLU
 

Action Decoder (RKN,LSTM): 1 fully connected + linear output:

    
  * Fully Connected 1: 512, ReLU
 
 
Transition Model (ac-RKN,RKN):
bandwidth: 3, number of basis: 15

    
  * α(z⃗_t): No hidden layers - softmax output


Control Model (ac-RKN):
1 fully connected + linear output

    
  * Fully Connected 1: 45, ReLU
 

Architecture For FFNN Baseline
3 fully connected + linear output

    
  * Fully Connected 1: 500, ReLU
    
  * Fully Connected 2: 500, ReLU
    
  * Fully Connected 3: 500, ReLU


Dropout Regularization - 0.563

Learning Rate - 1e-5

Optimizer Used: SGD Optimizer



§ HYPERPARAMETERS: CHAPTER <REF>




 §.§ Pneumatic Musculoskeletal Robot Arm

Recurrent Models


Context Encoder (HiP-RSSM): 1 fully connected + linear output  (elu + 1) 

    
  * Fully Connected 1: 240, ReLU
 
Observation Encoder (HiP-RSSM,RKN,LSTM,GRU): 1 fully connected + linear output  (elu + 1) 

    
  * Fully Connected 1: 120, ReLU
 
Observation Decoder (HiP-RSSM,RKN,LSTM): 1 fully connected + linear output:

    
  * Fully Connected 1: 120, ReLU
 
 
Transition Model (HiP-RSSM,RKN):
number of basis: 15 

    
  * α(z_t): No hidden layers - softmax output

Control Model (HiP-RSSM,RKN):
3 fully connected + linear output

    
  * Fully Connected 1: 120, ReLU
    
  * Fully Connected 2: 120, ReLU
    
  * Fully Connected 3: 120, ReLU
 



 §.§ Franka Robot Arm With Varying Loads

Recurrent Models


Encoder (HiP-RSSM,RKN,LSTM,GRU): 1 fully connected + linear output  (elu + 1) 

    
  * Fully Connected 1: 30, ReLU
 
Observation Decoder (HiP-RSSM,RKN,LSTM,GRU): 1 fully connected + linear output:

    
  * Fully Connected 1: 30, ReLU
 
 
Transition Model (HiP-RSSM,RKN):
number of basis: 32

    
  * α(z_t): No hidden layers - softmax output

Control Model (HiP-RSSM, RKN):
1 fully connected + linear output

    
  * Fully Connected 1: 120, ReLU
 



 §.§ Wheeled Robot Traversing Slopes Of Different Height

Recurrent Models


Encoder (HiP-RSSM,RKN,LSTM,GRU): 1 fully connected + linear output  (elu + 1) 

    
  * Fully Connected 1: 120, ReLU
 
Observation Decoder (HiP-RSSM,RKN,LSTM,GRU): 1 fully connected + linear output:

    
  * Fully Connected 1: 240, ReLU
 
 
Transition Model (HiP-RSSM,RKN):
number of basis: 15

    
  * α(z_t): No hidden layers - softmax output

Control Model (HiP-RSSM, RKN):
3 fully connected + linear output

    
  * Fully Connected 1: 120, ReLU
 




§ HYPERPARAMETERS: CHAPTER <REF>





  
Compute ResourcesFor training MTS3, LSTM, GRU and Transformer models we used compute nodes with (i) Nvidia 3090 and (ii) Nvidia 2080 RTX GPUs.
For training more computationally expensive locally linear models like RKN, HiP-RSSM we used
compute nodes with NVIDIA A100-40 GPUs.



  
Hyperparameters Hyperparameters were selected via grid search. In general, the performance of MTS3 is not very sensitive to hyperparameters. Among all the baselines, Transformer models were most sensitive to hyperparameters (see Appendix E.5 for details of Transformer architecture). 



  
Discretization Step: For MTS3, the discretization step for the slow time scale SSM as discussed in Section 3.1 for all datasets was fixed as H ·Δ t = 0.3 seconds. In our experiments, we found that discretization values between 0.2 ≤ H ·Δ t ≤ 0.5 seconds give similar performance.



  
Rule Of thumb for choosing discretization step in MTS3: For any N-level MTS3 as defined in Section 3.4, we recommend searching for discretization factor H_i as a hyperparameter. However, as a general rule of thumb, it can be chosen as H_i=(√(T))^i, where T is the maximum prediction horizon required / episode length. This ensures that very long recurrences are divided between smaller equal-length task-reconfigurable local SSM windows (of length √(T)) spread across several hierarchies.



  
Encoder Decoder Architecture: For all recurrent models (MTS3, HiP-RSSM, RKN, LSTM and GRU) we use a similar encoder-decoder architecture across datasets. Small variations from these encoder-decoder architecture hyperparameters can still lead to similar prediction performance as reported in the paper.


Observation Set Encoder (MTS3): 1 fully connected + linear output:

	
  * Fully Connected 1: 240, ReLU
 

Action Set Encoder (MTS3): 1 fully connected + linear output:

	
  * Fully Connected 1: 240, ReLU
 

Observation Encoder (MTS3, HiP-RSSM, RKN, LSTM, GRU): 1 fully connected + linear output:

	
  * Fully Connected 1: 120, ReLU
 

Observation Decoder (MTS3, HiP-RSSM, RKN, LSTM, GRU): 1 fully connected + linear output:

	
  * Fully Connected 1: 120, ReLU
 

Control Model (Primitive Action Encoder) (MTS3, HiP-RSSM, RKN):
1 fully connected + linear output:

	
  * Fully Connected 1: 120, ReLU
 

The rest of the hyperparameters are described below:

 



 §.§ D4RL Datasets




  §.§.§ Half Cheetah

Recurrent Models

 
Transition Model (HiP-RSSM, RKN):
number of basis: 32 


    
  * α(z_t): No hidden layers - softmax output


Autoregressive Transformer Baseline

Learning Rate: 1e-5

Optimizer Used: Adam Optimizer 

Embedding size: 96 

Number of Decoder Layers: 4 

Number Of Attention Heads: 4


Multistep Transformer Baseline

Learning Rate: 1e-5

Optimizer Used: Adam Optimizer 

Embedding size: 128 

Number Of Encoder Layers: 2

Number of Decoder Layers: 1 

Number Of Attention Heads: 4




  §.§.§ Franka Kitchen

Recurrent Models

 
Transition Model (HiP-RSSM, RKN):
number of basis: 15 


    
  * α(z_t): No hidden layers - softmax output


Autoregressive Transformer Baseline

Learning Rate: 5e-5

Optimizer Used: Adam Optimizer 

Embedding size: 64 

Number of Decoder Layers: 4 

Number Of Attention Heads: 4







Multistep Transformer Baseline

Learning Rate: 1e-5

Optimizer Used: Adam Optimizer 

Embedding size: 64 

Number Of Encoder Layers: 2

Number of Decoder Layers: 1 

Number Of Attention Heads: 4









  §.§.§ Maze 2D

Recurrent Models

 
Transition Model (HiP-RSSM, RKN):
number of basis: 15 


    
  * α(z_t): No hidden layers - softmax output


Autoregressive Transformer Baseline

Learning Rate: 5e-5

Optimizer Used: Adam Optimizer 

Embedding size: 96 

Number of Decoder Layers: 4 

Number Of Attention Heads: 4







Multistep Transformer Baseline

Learning Rate: 1e-5

Optimizer Used: Adam Optimizer 

Embedding size: 128 

Number Of Encoder Layers: 2

Number of Decoder Layers: 1 

Number Of Attention Heads: 4










 §.§ Franka Robot Arm With Varying Loads

Recurrent Models



 
Transition Model (HiP-RSSM,RKN):
number of basis: 32

    
  * α(z_t): No hidden layers - softmax output



Autoregressive Transformer Baseline

Learning Rate: 5e-5

Optimizer Used: Adam Optimizer 

Embedding size: 64 

Number of Decoder Layers: 4 

Number Of Attention Heads: 4







Multistep Transformer Baseline

Learning Rate: 2e-5

Optimizer Used: Adam Optimizer 

Embedding size: 64 

Number Of Encoder Layers: 2

Number of Decoder Layers: 1 

Number Of Attention Heads: 4









 §.§ Hydraulic Excavator



 
Transition Model (HiP-RSSM,RKN):
number of basis: 15

    
  * coefficient net α(z_t): No hidden layers - softmax output



Autoregressive Transformer Baseline

Learning Rate: 1e-5

Optimizer Used: Adam Optimizer 

Embedding size: 96 

Number of Decoder Layers: 4 

Number Of Attention Heads: 4







Multistep Transformer Baseline

Learning Rate: 5e-5

Optimizer Used: Adam Optimizer 

Embedding size: 64 

Number Of Encoder Layers: 2

Number of Decoder Layers: 1 

Number Of Attention Heads: 4












 §.§ Wheeled Robot Traversing Uneven Terrain


 
Transition Model (HiP-RSSM,RKN):
number of basis: 15

    
  * coefficient net α(z_t): No hidden layers - softmax output


Autoregressive Transformer Baseline

Learning Rate: 5e-5

Optimizer Used: Adam Optimizer 

Embedding size: 128 

Number of Decoder Layers: 4 

Number Of Attention Heads: 4







Multistep Transformer Baseline

Learning Rate: 5e-5

Optimizer Used: Adam Optimizer 

Embedding size: 64 

Number Of Encoder Layers: 4

Number of Decoder Layers: 2 

Number Of Attention Heads: 4









 §.§ Transformer Architecture Details

For the AR-Transformer Baseline, we use a GPT-like autoregressive version of transformers except that for the autoregressive input we also concatenate the actions to make action conditional predictions.

For Multi-Transformer we use the same direct multistep prediction and loss as in recent Transformer time-series forecasting literature <cit.>. A description of the action conditional direct multi-step version of the transformer is given in Algorithm <ref>.



