SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yu-Ren Liu\*,1,4, Biwei Huang \*,2, Zhengmao Zhu\*1,4, Honglong Tian\*1, Mingming Gong\*5,4, Yang Yu\*1,6,7, Kun Zhang\*3,4
$^{1}$  National Key Laboratory for Novel Software Technology, Nanjing University, China
<sup>2</sup> University of California San Diego, USA
3 Carnegie Mellon University, USA
<sup>4</sup> Mohamed bin Zayed University of Artificial Intelligence, UAE
<sup>5</sup> University of Melbourne, Australia
6 Polixir.ai, China
7 Peng Cheng Laboratory, China
{liuyr,zhuzm,tianhl,yuy}@lamda.nju.edu.cn bih007@ucsd.edu
mingming.gong@unimelb.edu.au kunz1@cmu.edu
Extracting a stable and compact representation of the environment is crucial for efficient reinforcement learning in high-dimensional, noisy, and non-stationary environments. Different categories of information coexist in such environments – how to effectively extract and disentangle the information remains a challenging problem. In this paper, we propose IFactor, a general framework to model four distinct categories of latent state variables that capture various aspects of information within the RL system, based on their interactions with actions and rewards. Our analysis establishes block-wise identifiability of these latent variables, which not only provides a stable and compact representation but also discloses that all reward-relevant factors are significant for policy learning. We further present a practical approach to learning the world model with identifiable blocks, ensuring the removal of redundancies but retaining minimal and sufficient information for policy optimization. Experiments in synthetic worlds demonstrate that our method accurately identifies the ground-truth latent variables, substantiating our theoretical findings. Moreover, experiments in variants of the DeepMind Control Suite and RoboDesk showcase the superior performance of our approach over baselines.
Humans excel at extracting various categories of information from complex environments [1, 2]. By effectively distinguishing between task-relevant information and noise, humans can learn efficiently and avoid distractions. Similarly, in the context of reinforcement learning, it's crucial for an agent to precisely extract information from high-dimensional, noisy, and non-stationary environments.
World models [3, 4] tackle this challenge by learning compact representations from images and modeling dynamics using low-dimensional features. Recent research has demonstrated that learning policies through latent imagination in world models significantly enhances sample efficiency [5, 6, 7, 8]. However, these approaches often treat all information as an undifferentiated whole, leaving policies susceptible to irrelevant distractions [9] and lacking transparency in decision-making [10].
This paper investigates the extraction and disentanglement of diverse information types within an environment. To address this, we tackle two fundamental questions: 1) How can we establish a
*: Equal contribution. †: Corresponding authors.
(a) Car-driving task.
Figure 1: (a) An illustrative example of the car-driving task. (b) Four categories of latent state variables in the car driving task. (c) The structure of our world model. Grey nodes denote observed variables and other nodes are unobserved. We allow causally-related latent processes for four types of latent variables and prove that they are respectively identifiable up to block-wise invertible transformations. We show that both  $s_t^{ar}$  and  $s_t^{\bar{a}r}$  are essential for policy optimization. The inclusion of the gray dashed line from  $(s_t^{ar}, s_t^{\bar{a}r})$  to  $a_t$  signifies that the action could be determined by reward-relevant variables. (d) The structure of Denoised MDP [9]. It assumes the latent processes of  $x_t$  and  $y_t$  are independent and uses only  $x_t$  for policy optimization. The gray dashed line from  $x_{t-1}$  to  $a_{t-1}$  shows that the action could be determined only based on the controllable and reward-relevant latent variables. Further, the existence of instantaneous causal effect from  $x_t$  and  $y_t$  to  $z_t$  renders the latent process unidentifiable without extra intervention on latent states [14].
comprehensive classification system for different information categories in diverse decision scenarios? 2) Can the latent state variables, classified according to this system, be accurately identified? Prior research has made progress in answering the first question. Task Informed Abstractions [11] partition the state space into reward-relevant and reward-irrelevant features, assuming independent latent processes for each category. Iso-Dream [12] learns controllable and noncontrollable sources of spatiotemporal changes on isolated state transition branches. Denoised MDP [9] further decomposes reward-relevant states into controllable and uncontrollable components, also assuming independent latent processes. However, the assumption of independent latent processes is overly restrictive and may lead to decomposition degradation in many scenarios. Moreover, none of these approaches guarantee the identifiability of representations, potentially leading to inaccurate recovery of the underlying latent variables. While discussions on the identifiability of representations in reinforcement learning (RL) exist under linear assumptions [13], the question of identifiability remains unexplored for general nonlinear cases.
In this paper, we present IFactor, a general framework to model four distinct categories of latent state variables within the RL system. These variables capture different aspects of information based on their interactions with actions and rewards, providing transparent representations of the following aspects: (i) reward-relevant and controllable parts, (ii) reward-irrelevant but uncontrollable parts, (iii) controllable but reward-irrelevant parts, and (iv) unrelated noise (see Section 2). Diverging from prior methods, our approach employs a general factorization for the latent state variables, allowing for causally-related latent processes. We theoretically establish the block-wise identifiability of these four variable categories in general nonlinear cases under weak and realistic assumptions. Our findings challenge the conclusion drawn in the Denoised MDP [9] that only controllable and reward-relevant variables are required for policy optimization. We emphasize the necessity of considering states that directly or indirectly influence the reward during the decision-making process, irrespective of their controllability. To learn these four categories of latent variables, we propose a principled approach that involves optimizing an evidence lower bound and integrating multiple novel mutual information constraints. Through simulations on synthetic data, we demonstrate the accurate identification of true latent variables, validating our theoretical findings. Furthermore, our method achieves state-of-the-art performance on variants of RoboDesk and DeepMind Control Suite.
2 Four Categories of (Latent) State Variables in RL
For generality, we consider tasks in the form of Partially Observed Markov Decision Process (POMDP) [15], which is described as  $\mathcal{M} \triangleq (S, \mathcal{A}, \Omega, R, T, O, \gamma)$ , where  $S$  is the latent state space,  $\mathcal{A}$  is the action space,  $\Omega$  is the observation space,  $R: S \times A \to \mathbb{R}$  defines the reward function,  $T: S \times A \to \Delta(S)$  is the transition dynamics,  $O: S \times A \to \Delta(\Omega)$  is the observation function,  $\gamma \in [0, 1]$  is the discount factor. We use  $\Delta(S)$  to denote the set of all distributions over  $S$ . The agent can interact with the environment to get sequences of observations  $\{\langle o_t, a_t, r_t \rangle\}_{t=1}^T$ . The objective is to find a policy acting based on history observations, that maximizes the expected cumulative (discounted) reward.
In the context of POMDPs, extracting a low-dimensional state representation from high-dimensional observations is crucial. Considering that action and reward information is fundamental to decision-making across various scenarios, we disentangle the latent variables in the environment into four distinct categories, represented as  $\mathbf{s}_{\mathbf{t}} = \{s_{t}^{ar}, s_{t}^{\bar{a}r}, s_{t}^{a\bar{r}}, s_{t}^{\bar{a}\bar{r}}\}$ , based on their relationships with action and reward (see Figure 1(c) as a graphical illustration):
- Type 1:  $s_t^{ar}$  has an incident edge from  $a_{t - 1}$ , and there is a directed path from  $s_t^{ar}$  to  $r_t$ .
- Type 2:  $s_t^{\bar{a}r}$  has no incident edge from  $a_{t-1}$ , and there is a directed path from  $s_t^{\bar{a}r}$  to  $r_t$ .
- Type 3:  $s_t^{a\bar{r}}$  has an incident edge from  $a_{t - 1}$ , and there is no directed path from  $s_t^{a\bar{r}}$  to  $r_t$ .
- Type 4:  $s_t^{\bar{a}\bar{r}}$  has no incident edge from  $a_{t - 1}$ , and there is no directed path from  $s_t^{\bar{a}\bar{r}}$  to  $r_t$ .
$s_t^{ar}$  represents controllable and reward-relevant state variables that are essential in various scenarios. Taking the example of a car driving context,  $s_t^{ar}$  encompasses driving-related states like the current speed, position, and direction of the car. These latent variables play a critical role in determining the driving policy since actions such as steering and braking directly influence these states, which, in turn, have a direct impact on the reward received.
$s_{t}^{\bar{a}r}$  refers to reward-relevant state variables that are beyond our control. Despite not being directly controllable, these variables are still necessary for policy learning. In the context of car driving,  $s_{t}^{\bar{a}r}$  includes factors like surrounding vehicles and weather conditions. Although we cannot directly control other cars, our car's status influences their behavior: if we attempt to cut in line unethically, the surrounding vehicles must decide whether to yield or block our behavior. As a result, we need to adjust our actions based on their reactions. Similarly, the driver must adapt their driving behavior to accommodate various weather conditions, even though the weather itself is uncontrollable.
The state variables denoted as  $s_t^{a\bar{r}}$  consist of controllable but reward-irrelevant factors. Examples of  $s_t^{a\bar{r}}$  could be the choice of music being played or the positioning of ornaments within the car. On the other hand,  $s_t^{\bar{a}\bar{r}}$  represents uncontrollable and reward-irrelevant latent variables such as the remote scenery. Both  $s_t^{a\bar{r}}$  and  $s_t^{\bar{a}\bar{r}}$  do not have any impact on current or future rewards and are unrelated to policy optimization.
We next build a connection between the graph structure and statistical dependence between the variables in the RL system, so that the different types of state variables can be characterized from the data. To simplify symbol notation, we define  $s_t^r \coloneqq (s_t^{ar}, s_t^{\bar{a} r})$ ,  $s_t^{\bar{r}} \coloneqq (s_t^{\bar{a}\bar{r}}, s_t^{\bar{a}\bar{r}})$ ,  $s_t^a \coloneqq (s_t^{a\bar{r}}, s_t^{a\bar{r}})$  and  $s_t^{\bar{a}} \coloneqq (s_t^{\bar{a}\bar{r}}, s_t^{\bar{a}\bar{r}})$ . Specifically, the following proposition shows that  $s_t^r$  that has directed paths to  $r_{t + \tau}$  (for  $\tau > 0$ ), is minimally sufficient for policy learning that aims to maximize the future reward and can be characterized by conditional dependence with the cumulative reward variable  $R_{t} = \sum_{t' = t} \gamma^{t' - t} r_{t'}$  which has also been shown in [13].
Proposition 1. Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful to the measured data,  $s_t^r \subseteq \mathbf{s_t}$  is a minimal subset of state dimensions that are sufficient for policy learning, and  $s_{i,t} \in s_t^r$  if and only if  $s_{i,t} \nsubseteq R_t|a_{t - 1:t}, s_{t - 1}^r$ .
Moreover, the proposition below shows that  $s_t^a$ , that has a directed edge from  $a_{t-1}$ , can be directly controlled by actions and can be characterized by conditional dependence with the action variable.
Proposition 2. Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful to the measured data,  $s_t^a \subseteq \mathbf{s_t}$  is a minimal subset of state dimensions that are sufficient for direct control, and  $s_{i,t} \in s_t^a$  if and only if  $s_{i,t} \nleq a_{t-1} | \mathbf{s}_{t-1}$ .
Furthermore, based on Proposition 1 and Proposition 2, we can further differentiate  $s_t^{ar}$ ,  $s_t^{\bar{a}\bar{r}}$ ,  $s_t^{\bar{a}\bar{r}}$  from  $s_t^r$  and  $s_t^a$ , which is given in the following proposition.
Proposition 3. Under the assumption that the graphical representation, corresponding to the environment model, is Markov and faithful to the measured data, we can build a connection between the graph structure and statistical independence of causal variables in the RL system, with (1)  $s_{i,t} \in s_t^{ar}$  if and only if  $s_{i,t} \not\perp R_t|a_{t-1:t}, s_{t-1}^r$  and  $s_{i,t} \not\perp a_{t-1}|s_{t-1}$ , (2)  $s_{i,t} \in s_t^{\bar{a}r}$  if and only if  $s_{i,t} \not\perp R_t|a_{t-1:t}, s_{t-1}^r$  and  $s_{i,t} \not\perp a_{t-1}|s_{t-1}$ , (3)  $s_{i,t} \in s_t^{a\bar{r}}$  if and only if  $s_{i,t} \not\perp R_t|a_{t-1:t}, s_{t-1}^r$  and  $s_{i,t} \not\perp a_{t-1}|s_{t-1}$ , and (4)  $s_{i,t} \in s_t^{\bar{a}\bar{r}}$  if and only if  $s_{i,t} \not\perp R_t|a_{t-1:t}, s_{t-1}^r$  and  $s_{i,t} \not\perp a_{t-1}|s_{t-1}$ .
By identifying these four categories of latent state variables, we can achieve interpretable input for policies, including (i)reward-relevant and directly controllable parts, (ii) reward-relevant but not controllable parts, (iii) controllable but non-reward-relevant parts, and (iv) unrelated noise. Furthermore, policy training will become more sample efficient and robust to task-irrelevant changes by utilizing only  $s_t^{ar}$  and  $s_t^{\bar{a}\bar{r}}$  for policy learning. We show the block-wise identifiability of the four categories of variables in the next section.
3 Identifiability Theory
Identifying causally-related latent variables from observations is particularly challenging, as latent variables are generally not uniquely recoverable [16, 17]. Previous work in causal representation learning and nonlinear-ICA has established identifiability conditions for non-parametric temporally latent causal processes [18, 19]. However, these conditions are often too stringent in reality, and none of these methods have been applied to complex control tasks. In contrast to component-wise identifiability, our focus lies on the block-wise identifiability of the four categories of latent state variables. From a policy optimization standpoint, this is sufficient because we do not need the latent variables to be recovered up to permutation and component-wise invertible nonlinearities. This relaxation makes the conditions more likely to hold true and applicable to a wide range of RL tasks. We provide proof of the block-wise identifiability of the four types of latent variables. To the best of our knowledge, we are the first to prove the identifiability of disentangled latent state variables in general nonlinear cases for RL tasks.
According to the causal process in the RL system (as described in Eq.1 in [13]), we can build the following mapping from latent state variables  $\mathbf{s}_t$  to observed variables  $o_t$  and future cumulative reward  $R_t$ :
$$
\left[ o _ {t}, R _ {t} \right] = f \left(s _ {t} ^ {r}, s _ {t} ^ {\bar {r}}, \eta_ {t}\right), \tag {1}
$$
$$
\begin{array}{r c l} o _ {t} & = & f _ {1} \left(s _ {t} ^ {r}, s _ {t} ^ {\bar {r}}\right), \\ R _ {t} & = & f _ {2} \left(s _ {t} ^ {r}, \eta_ {t}\right). \end{array} \tag {2}
$$
Here, note that to recover  $s_t^r$ , it is essential to take into account all future rewards  $r_{t:T}$ , because any state dimension  $s_{i,t} \in \mathbf{s}_t$  that has a directed path to the future reward  $r_{t + \tau}$ , for  $\tau > 0$ , is involved in  $s_t^r$ . Hence, we consider the mapping from  $s_t^r$  to the future cumulative reward  $R_{t}$ , and  $\eta_t$  represents residuals, except  $s_t^r$ , that have an effect to  $R_{t}$ .
Below, we first provide the definition of blockwise identifiability and relevant notations, related to [20, 21, 22, 23].
Definition 1 (Blockwise Identifiability). A latent variable  $\mathbf{s}_{\mathbf{t}}$  is blockwise identifiable if there exists a one-to-one mapping  $h(\cdot)$  between  $\mathbf{s}_{\mathbf{t}}$  and the estimated  $\hat{\mathbf{s}}_{\mathbf{t}}$ , i.e.,  $\hat{\mathbf{s}}_{\mathbf{t}} = h(\mathbf{s}_{\mathbf{t}})$ .
Notations. We denote by  $\tilde{\mathbf{s}}_t\coloneqq (s_t^r,s_t^{\bar{r}},\eta_t)$  and by  $|s|$  the dimension of a variable  $s$ . We further denote  $d_{s_r}\coloneqq |s_r^r |,d_{s_{\bar{r}}}\coloneqq |s_{\bar{r}}^{\bar{r}}|,d_{\bar{s}}\coloneqq |\tilde{\mathbf{s}}_t|,d_o\coloneqq |o_t|$ , and  $d_R\coloneqq |R_t|$ . We denote by  $\mathcal{F}$  the support of Jacobian  $\mathbf{J}_f(\mathbf{s}_t)$ , by  $\hat{\mathcal{F}}$  the support of  $\mathbf{J}_{\hat{f}}(\hat{\mathbf{s}}_t)$ , and by  $\mathcal{T}$  the support of  $\mathbf{T}(\mathbf{s}_t)$  with  $\mathbf{J}_{\hat{f}}(\hat{\mathbf{s}}_t) = \mathbf{J}_f(\mathbf{s}_t)\mathbf{T}(\mathbf{s})$ . We also denote  $T$  as a matrix with the same support as  $\mathcal{T}$ .
In addition, given a subset  $\mathcal{D} \subseteq \{1, \dots, d_{\tilde{s}}\}$ , the subspace  $\mathbb{R}_{\mathcal{D}}^{d_{\tilde{s}}}$  is defined as  $\mathbb{R}_{\mathcal{D}}^{d_{\tilde{s}}} \coloneqq \{z \in \mathbb{R}^{d_{\tilde{s}}} | i \notin \mathcal{D} \Rightarrow z_i = 0\}$ . In other words,  $\mathbb{R}_{\mathcal{D}}^{d_{\tilde{s}}}$  refers to the subspace of  $\mathbb{R}^{d_{\tilde{s}}}$  indicated by an index set  $\mathcal{D}$ .
We next show that the different types of states  $s_t^{ar}$ ,  $s_t^{\bar{a}r}$ ,  $s_t^{a\tilde{r}}$ , and  $s_t^{\tilde{a}\tilde{r}}$  are blockwise identifiable from observed image variable  $o_t$ , reward variable  $r_t$ , and action variable  $a_t$ , under reasonable and weak assumptions, which is partly inspired by [20, 21, 22, 23].
Theorem 1. Suppose that the causal process in the RL system and the four categories of latent state variables can be described as that in Section 2 and illustrated in Figure 1(c). Under the following assumptions
A1. The mapping  $f$  in Eq. 1 is smooth and invertible with smooth inverse.
A2. For all  $i \in \{1, \ldots, d_{o} + d_{R}\}$  and  $j \in \mathcal{F}_{i,:}$ , there exist  $\{\tilde{\mathbf{s}}_{t}^{(l)}\}_{l=1}^{|F_{i,:}|}$  so that  $\text{span} \{\mathbf{J}_{f}(\tilde{\mathbf{s}}_{t}^{(l)})_{i,:}\}_{l=1}^{|F_{i,:}|} = \mathbb{R}_{\mathcal{F}_{i,:}}^{d_{\tilde{s}}}$ , and there exists a matrix  $T$  with its support identical to that of  $\mathbf{J}_{\hat{f}}^{-1}(\hat{\mathbf{s}}_{t})\mathbf{J}_{f}(\tilde{\mathbf{s}}_{t})$ , so that  $[\mathbf{J}_{f}(\tilde{\mathbf{s}}_{t}^{(l)})T]_{j,:} \in \mathbb{R}_{\hat{\mathcal{F}}_{i,:}}^{d_{\tilde{s}}}$ .
Then, reward-relevant and controllable states  $s_t^{ar}$ , reward-relevant but not controllable states  $s_t^{\bar{a}\bar{r}}$ , reward-irrelevant but controllable states  $s_t^{\bar{a}\bar{r}}$ , and noise  $s_t^{\bar{a}\bar{r}}$ , are blockwise identifiable.
In the theorem presented above, Assumption A1 only assumes the invertibility of function  $f$ , while functions  $f_{1}$  and  $f_{2}$  are considered general and not necessarily invertible, as that in [23]. Since the function  $f$  is the mapping from all (latent) variables, including noise factors, that influence the observed variables, the invertibility assumption holds reasonably. However, note that it is not
reasonable to assume the invertibility of the function  $f_{2}$  since usually, the reward function is not invertible. Intuitively, Assumption A2 requires that the Jacobian varies "enough" so that it cannot be contained in a proper subspace of  $\mathbb{R}_{\mathcal{F}_{i:}}^{d_s}$ . This requirement is necessary to avoid undesirable situations where the problem becomes ill-posed and is essential for identifiability. A special case when this property does not hold is when the function  $f$  is linear, as the Jacobian remains constant in such cases. Some proof techniques of Theorem 1 follow from [20, 22, 23], with the detailed proof given in Appendix A.4.
4 World Model with Disentangled Latent Dynamics
Based on the four categories of latent state variables, we formulate a world model with disentangled latent dynamics. Each component of the world model is described as follows:
$$
\left\{ \begin{array}{l l} \text {O b s e r v a t i o n M o d e l :} & p _ {\theta} \left(o _ {t} \mid \mathbf {s} _ {\mathbf {t}}\right) \\ \text {R e w a r d M o d e l :} & p _ {\theta} \left(r _ {t} \mid s _ {t} ^ {r}\right) \\ \text {T r a n s i t i o n M o d e l :} & p _ {\gamma} \left(\mathbf {s} _ {\mathbf {t}} \mid \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}\right) \\ \text {R e p r e s e n t a t i o n M o d e l :} & q _ {\phi} \left(\mathbf {s} _ {\mathbf {t}} \mid o _ {t}, \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}\right) \end{array} \right. \tag {3}
$$
Specifically, the world model includes three generative models: an observation model, a reward function, and a transition model (prior), and a representation model (posterior). The observation model and reward model are parameterized by  $\theta$ . The transition model is parameterized by  $\gamma$  and the representation model is parameterized by  $\phi$ . We assume noises are i.i.d for all models. The action  $a_{t-1}$  has a direct effect on the latent states  $\mathbf{s}_{t}$ , but not on the perceived signals  $o_t$ . The perceived signals,  $o_t$ , are generated from the underlying states  $\mathbf{s}_{t}$ , while signals  $r_t$  are generated only from reward-relevant latent variables  $s_t^r$ .
According to the graphical model depicted in Figure 1(c), the transition model and the representation model can be further decomposed into four distinct sub-models, according to the four categories of latent state variables, as shown in equation 4.
Disentangled Transition Model: Disentangled Representation Model:
$$
\left\{ \begin{array}{l} p _ {\gamma_ {1}} \left(s _ {t} ^ {a r} \mid s _ {t - 1} ^ {r}, a _ {t - 1}\right) \\ p _ {\gamma_ {2}} \left(s _ {t} ^ {\bar {a} r} \mid s _ {t - 1} ^ {r}\right) \\ p _ {\gamma_ {3}} \left(s _ {t} ^ {a \bar {r}} \mid \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}\right) \\ p _ {\gamma_ {4}} \left(s _ {t} ^ {\bar {a} \bar {r}} \mid \mathbf {s} _ {\mathbf {t} - 1}\right) \end{array} \right. \quad \left\{ \begin{array}{l} q _ {\phi_ {1}} \left(s _ {t} ^ {a r} \mid o _ {t}, s _ {t - 1} ^ {r}, a _ {t - 1}\right) \\ q _ {\phi_ {2}} \left(s _ {t} ^ {\bar {a} r} \mid o _ {t}, s _ {t - 1} ^ {r}\right) \\ q _ {\phi_ {3}} \left(s _ {t} ^ {a \bar {r}} \mid o _ {t}, \mathbf {s} _ {\mathbf {t} - 1}, a _ {t - 1}\right) \\ q _ {\phi_ {4}} \left(s _ {t} ^ {\bar {a} \bar {r}} \mid o _ {t}, \mathbf {s} _ {\mathbf {t} - 1}\right) \end{array} \right. \tag {4}
$$
Specifically, we have  $p_{\gamma} = p_{\gamma_1} \cdot p_{\gamma_2} \cdot p_{\gamma_3} \cdot p_{\gamma_4}$  and  $q_{\phi} = q_{\phi_1} \cdot q_{\phi_2} \cdot q_{\phi_3} \cdot q_{\phi_4}$ . Note this factorization differs from previous works [11, 9] that assume independent latent processes. Instead, we only assume conditional independence among the four categories of latent variables given  $\mathbf{s}_{t-1}$ , which provides a more general factorization. In particular, the dynamics of  $s_t^{ar}$  and  $s_t^{\bar{a}\bar{r}}$  are dependent on  $s_{t-1}^r$ , ensuring that they are unaffected by any reward-irrelevant latent variables present in  $s_{t-1}^{\bar{r}}$ . On the other hand,  $s_t^{\bar{r}}$  may be influenced by all the latent variables from the previous time step. Note that the connections between  $\mathbf{s}_{t-1}$  and  $s_t^{\bar{r}}$  can be adjusted based on the concrete problem. Since  $s_t^{ar}$  and  $s_t^{a\bar{r}}$  are controllable variables, their determination also relies on  $a_{t-1}$ .
4.1 World Model Estimation