Guaranteed Discovery of Control-Endogenous Latent States with Multi-Step Inverse Models
Authors: Alex Lamb
Riashat Islam
Yonathan Efroni
Aniket Rajiv Didolkar
Dipendra Misra
Dylan J Foster
Lekan P Molu
Rajan Chari
Akshay Krishnamurthy
John Langford
Venue: Accepted by TMLR ()

Abstract: In many sequential decision-making tasks, the agent is not able to model the full complexity of the world, which consists of multitudes of relevant and irrelevant information. For example, a person walking along a city street who tries to model all aspects of the world would quickly be overwhelmed by a multitude of shops, cars, and people moving in and out of view, each following their own complex and inscrutable dynamics.  Is it possible to turn the agent's firehose of sensory information into a minimal latent state that is both necessary and sufficient for an agent to successfully act in the world? We formulate this question concretely, and propose the Agent Control-Endogenous State Discovery algorithm (AC-State), which has theoretical guarantees and is practically demonstrated to discover the minimal control-endogenous latent state which contains all of the information necessary for controlling the agent, while fully discarding all irrelevant information.    This algorithm consists of a multi-step inverse model (predicting actions from distant observations) with an information bottleneck.  AC-State enables localization, exploration, and navigation without reward or demonstrations.  We demonstrate the discovery of the control-endogenous latent state in three domains: localizing a robot arm with distractions (e.g., changing lighting conditions and background), exploring a maze alongside other agents, and navigating in the Matterport house simulator.  

Under review as submission to TMLR
Guaranteed Discovery of Control-Endogenous Latent States with Multi-Step Inverse Models
Anonymous authors Paper under double-blind review
Abstract
In many sequential deeision-makee making tasks, the agent is not able model the full complex- ity of the world, which consists of multitudes of relevant and irrelevant information. For example, a person walking along a city street who tries to model all aspects of the world would quickly be overwhelmed by a multitude of shops, cars, and people moving in and out of view, each following their own complex and inscrutable dynamics. it possible to turn the agent's firehose of sensory information into a minimal latent state that is both necessary and sufficient. for an agent successfully act the world? We formulate this question concretely, and propose the Agent Control-Endogenous State Discovery algorithm (AC-State), which has theoretical guarantees and is practically demonstrated to discover the minimal control-endogenon latent. slate which contains all of the information necessary for controlling the agent, while fully discarding all irrelevant information. This algorithm consists of a multi-step inverse model (predicting actions from distant observations) with an information bottleneck. AC-State enables localization, exploration, and navigation without reward or demonstrations. We demonstrate the discovery of the control-ondogenous latent state in three domains: localizing robot arm with distractions (e-g., changing lighting conditions and background), exploring a maze alongside other agents, and navigating in the Matterport house simulator.
1 Introduction
In many real-world systems, the observation space is generated by multiple complex and sparsely interacting subsystems. The state that is necessary for controlling an agent depends only on small fraction of the information in the observation space. For example, consider a world consisting of a robot arm that is controlled by an agent with access to a high-resolution vidco of the robot. The agent's observation intertwines information that the agent controls with exogenous content such as lighting conditions or videos in the background. We the only information that either can be controlled by the agent (such as an object on the table that can be manipulated by the arm) or affects the agent's control (c.g., obstacle blocking the robot arm's motion). Discovering this representation while ignoring irrelevant information offers the promise of vastly improved planning, exploration, and interpretability.
latend state as the parsimonious representation, which includes
A key challenge to discovering these representations in real world applications, is that only a continual stream of observations along with the agent's actions are available. For example, consider discovering the control- endogenous latent state from vidco of a robot arm along with the sequence of actions taken control the robot. Efroni et al. (2022c) demonstrated an algorithm for discovering control-endooeenus latent state that open-loop planning with the agent reset to a fixed start state at the beginning of each episode, which is impractical in most real domains. Prior works [Efroni et al. 2022a Wang et al. 2022b) also discovered the control-endogenous latent state by assuming access to a given factorized encoder, which is often not available in practice.
In this work, we consider a setting where the agent acts in world consisting of complex observations with relevant and irrelevant information. We consider environments where observations are generated from states that be decoupled into control-endogenous and ezogenous components. Our definition of exogenous

---

Under review as submission to TMLR
strictly refers to aspects of the world that never interact with the agent, which may include things like the detailed visual textures of objects background processes that the agent cannot. interact with. It is vital to discard irrelevant parts of the observation space. Our proposed algorithm predicts the first action that needs to be taken from any current observation to reach any future state a certain number of steps in the future. We establish some theoretical results showing that this captures a complete and minimal control-endogenous state. This has many appcaling properties in that we never need to learn a sequence model or generative model. We only need to use supervised learning predict actions. Prior works have considered partially observable settings to summarize agent-states (Littman et al. 1995 Kaelbling et al,, 1998), or quantify agent states based on predictions of rewards (Littman et al. 2001). We emphasize here the salience of information for controlling the agent in fully-observed settings.
a2
S
AC-State Discovers the smallest eonrol-endoognouss state assuming factorized dynamics
C
+
4
a,
4
AC-State collects data single random action followed by endogenous policy k-1 steps
NC (ftx,,,k) predict first action from x to reach x
x
14k
AC-State learns an encoder f for fix) by optimizing multi-step inverse model with bottleneck
L bonteauck (fxxt)
g Lasettrra ((xtt+
The control-endogenouu latent dynamics (top) can be discovered using AC-State, without learning model of the exogenous noise or the observation æ. This involves collecting data under a high-coverage endogenous policy (center) while optimizing a multi-step inverse model with a bottleneck (bottom).
Figure 1: Overview of AC-State
How can discover this control-endogenous latent state from raw observations and actions? We intro- duce the Agent Control-Endogenous (AC-State) algorithm, which provably guarantees discovery of the control-endogenoots latent state by excluding all aspects of the observations that are unnecessary for control. AC-State learns an encoder that maps given obscrvation to the corresponding control-cndogenous latent state (xx) This is accomplished by optimizing a novel objective using an expressive model class such as deep neural networks [Leuun et al.) 20115.. The proposed algorithm uses a multi-step inverse model as objective, along with a bottleneck on the capacity of the learned representation. we use a discrete information bottleneck, based on vector quantization, which plays a key role in discarding exogenous information.
we discuss later,
2

---

Under review as submission to TMLR
1-Step
AC-State
Algorithms
PPE
OSSR
DBC
CDL
Denoised-Mo
Inverse
(Ours)
Exogenous Invariant State
V
V
Exogenous Invariant Learning
X
X
Flexible Encoder
X
V
X
V
YOLO (No Resets) Setting
X
V
d
Reward Free
V
X
V
/
Control-Bndogenous Rep-
X
V
X
An Overview of the Properties of prior works on representation learning in RL is shown, with particular emphasis on robustness to exogenous information. We compare with several baselines including PPE Efroni et al. (2022c), OSSR Efroni et al.] (022a)) DBC Zhang et al. (2021a) Denoised MDP Wang et al. (022a)) and 1-Step Inverse Models Pathak et al. (2017). The comparison to AC-State aims to be as generous as possible the baselines. is used to indicate a known counterexample for given property.
Table 1:
In our experiments, we demonstrate the discovery of agent control-endooennuu latent states that are (nearly) identical to their ground-truth values, while only using access to the raw observations and actions. We measure the exactness of the correspondence between the ground truth latent state (not used for training) and the learned latent state. The AC-State approach requires no reward sigual, and many experiments require a few thousand interactions with the environment for initial learning along with zero additional examples to discover an appropriate policy, At the same time, successful discovery of the latent state implies the ability to localize, explore, and plan to reach goal states from single shared representation.
2 Prior Approaches
Deep learning architectures can be optimized for a wide range of differentiable objective functions. Our key question is: what is an objective for provably learning control-endogennus latent state that is compatible with deep learning? At issue is finding parsimonious representations that are sulficient for control of a dynamical system given observations from rich sensors (such as high-resolution videos) while discarding irrelevant details. Approaches such as optimal state estimation (Durrant- Whyte & Bailey 2006), system identification (Ljung 1998), and simultaneous localization and mapping (Cadena et al. 2016 Dissanayake et al. 2001) achieve parsimonious state estimation for control, yet require more domain expertise and design than desirable. Previous bearning-based approaches failed capture the full control-endogenous latent state exclude all irrelevant information (Efroni al. 2022c). Reinforcement learning approaches that
capture the latent state from rich sensors by employing autoencoders [Goodfellow et al. 2016) contrastive learning often capture noise components
Learning latent states for interactive environments is a mature research area with prolific contributions. We discuss few of the most important lines of research and how they fail to achieve guaranteed discovery of the control-cndogenous latent state. We categorize these contributions into broad areas based on what they predict: rewards (or value functions), future latent states, observations, the relationship between observa- tions, and/or actions. Table 1 provides summarized comparison of the properties of AC-State and prior works. In this table, algorithms that must learn the full exogenous state before learning how to discard it are marked as having an exogenous invariant state but not exogenous invariant learning.
Limitations of Predicting Rewards: Reward-based bisimulation (Zhang et al. 2021b) can filter irrelevant information from the latent state but dependent on access to a reward signal. Deep reinforcement learning based on reward optimization (Mnih al.) often struggles with sample complexity when the reward is sparse and fails completely when the reward is absent. An adversarial approach can also be used to learn disentangled reward-relovant. and roward-irrclevant state (uu ct al. 2021 This can successfully ignore
CConsider divided freeway, where cars travel on opposing sides of the lane. For this situation, autoencoders or contrastive learning objectives (commonly referred as "seltsuperrised learning") produce distinct latent states for every unique config- uration of cars the other side of the lane divider. For example, autoencoder or generative model would learn predict the full configuration and visual details of all the cars, even those that could interact with the agent.
3

---

Under review as submission to TMLR
distractors but may struggle if the distractors are difficult to completely model and can only be used when external reward signal is available.
Limitations of Predicting Latent States: Approaches that involve predicting future latent states from past latent states have achieved good performance, but there is no theoretical guarantee that the latent state captures the full control-endogennuu latent state (Guo et al. 2022 Schwarzer et al.] 2021bb Ye et al. 2021 Pathak et al. 2017). Deep bisimulation approaches learn state representations for control tasks with agnosticism toward task-irrelevant details (Zhang et al., 2021a). Prior approaches have involved model that predicts latent states such that the pre-trained representations can solve a task (Schwarzer et al, 2021a), yet this approach is not task-agnostic, and it is not guaranteed to recover the control-cndogenom state. An auto-encoder trained with reconstruction loss or a dynamics model (Lange et al. 2012 Wahlström le al. 2015 Watter et al.) 2015) learns low-dimensional state representations to capture information relevant for solving a task. Although learning latent state representations has been shown to be useful for solving tasks, there is no guarantee that such methods can fully recover the underlying canroo-cndooenous latent state.
Limitations of Predicting Observations: Other works have also used generative models autoencoders to predict future observations for learning latent reprosentations, mostly for purposes of exploration. The idea of predicting observations is often referred to as "intrinsic motivation" to guide the agent towards exploring unseen regions of state space (Oudeyer & Kaplan 2007). Other works use autoencoders to estimate future observations in the feature space for exploration (Stadie et al. 2015), though such models can also fail in the presence of exogenous observations. Dynamics models learn to predict distributions over future observations, but often for the purpose of planning sequence of actions instead of recovering latent states for purposes of exploration. Models which predict observations directly must also predict exogenous noise (Misra et al., 2020). Some approaches that learn generative models in the observation space attempt to learn a further decomposition of the dynamics into control-endogenons and exogenous latent states. While this can result in the correct latent state, it still requires learning a full genorative model over the entire latent state, not just. the control-endoogeouss latent state (Wang et al. 2022a b).
Limitations of Predicting Relationships between Observations: By learning to predict relations between two consccutive observations, prior works have attempted to learn latent state representations, both theoretically (Misra el al.
2020) and empirically (Mazoure et al. 2020). For example, by exploiting mutual information based objectives (information gain based on current states and actions with future 2012), previous works have attempted to learn control-endogenous
states) Mazoure et al. 2020 Song et al. states in the presence of exogenous noise. However, unlike AC-State, they learn latent states dependent on exogenous noise, even though the learned representation can be useful for solving complex tasks (Mazoure ett al. 2020). Theoretically, Misra et al. (2020) uses a contrastive loss based objective to provably learn latent state representations that can be useful for hard exploration tasks. However, contrastive loss based representations have a counterexample where they are forced to capture exogenous noise (Efroni et al, 2022c), whereas AC-State exploits an exogenous free rollout policy with multi-step inverse dynamics model to provably and experimentally recover the full control-endoogenous latent state.
Limitations of Predicting Actions: AC-State aims at recovering the endtrol-e latent states by training a multi-step inverse dynamics model in the presence of exogenous noise. While prior works have explored similar objectives, either for exploration or for learning state representations, they are unable to recover latent states with perfect accuracy in the presence of exogenous noise (Efroni et al. 2022c). Tn Appendix B we provide concrete counter-example demonstrating why a one-step inverse objective fails to learn the correct latent state. Approaches based on action prediction with one-step inverse models (Pathak et al. 2017) are widely used in practice (Baker et al. 2022 Badia et al. 2020). However, these inverse models can fail to capture the full condrol-endogenous latent dynamics (Eironi et al. 2022c Hutter & Hansen 2022), while combining them with an autoencoder (Bharadhwaj et al. 2022) inherits the weaknesses of that approach.
Limitations of Empowerment-basee objectives Empowerment-bassed objectives focus the idea that an agent should try to seek out states where it is empowered by having the greatest number of possible states that it can easily reach (Klyubin et al. 2005). For example, in a maze with two rooms, the most
4

---

Under review as submission to TMLR
empowered state is the doorway, since it makes it easy to reach either of the rooms. Concrete instantiations of the empowerment objective may involve training models to predict the distribution of actions from ob- servations (either single-step or multi-step inverse models) (Mohamed & Rezende 2015 Yu et al.) 2019).
but they lack the information bottleneck term and the requirement of an exogenous- exogenous-independen rollout policy. The analysis and theory in this work focus on action prediction as a particular method for measuring empowerment rather than as a way of guarantceing the discovery of a minimal control-endoenous latent state and ignoring exogenous noise.
Limitations of Causal Learning Approaches: The concept of control-endogenonn latent states is closely related to causality in the sense that it captures the content in the state space that is causally related to actions. This perspective has been explored by [Wang et al. (2022b), which used conditional independence testing to find sets of high-level variables that are related to actions. Time series with known variables have also been used to discover causal relationships (Mastakouri et al.) 2021). A pre-trained feature extractor for high-level variables (such as keypoints in computer vision) can also be used as the inputs for discovering causal structure (Li et. al., 2020). These methods assume access to a small set of high-level variables with sparse causal structure. This assumption allows these methods to completely avoid the representational learning aspect of the problem. Learning causal dynamics end-to-end using doep learning is nascent area of research, even when high-level variables are provided as the inputs for learning (Subramanian et al. 2022:
Scherrer et al. 2021).
Limitations of other approaches using Multi-Step Inverse Models: Here we discuss several recent works that study the reinforcement learning problem in the presence of exogenous and irrelevant information. Efroni et al. 2022cb formulated the Ex-BMDP model and designed a provably efficient algorithm that learns the state representation. However, their algorithm succeeds only in the episodic setting, when the initial cettrol-endogenous state is initialized deterministically. This is used alongside deterministic dynamics to use open-loop plans (i.e. plans which do not look at the actual state of the environment) to construct the endogenous latent state. These strict assumptions make their algorithm impractical in many cases of interest. Here we removed the deterministic assumption of the initial latent state. Indeed, as our theory suggests, AC-State can be applied in the "you-only- live-once" (YOLO) setting when an agent has access to single trajectory. In Efroni et al. (2022b) and OSSR (Efroni et al., 2022a) the authors designed provable RL algorithms that efficiently learn in the presence of exogenous noise under different assumptions about the underlying dynamics. These works, however, focused on statistical aspects of the problem; how to scale these approaches to complex environments and combine function approximations is currently unknown and scems challenging. Efroni et al. (2022c) proposes a deterministic path planning based algorithm, which may not be directly compatible with most existing deep RL frameworks. In contrast, the proposed AC-State algorithm recovers the underlying latent states and can be directly used to construct tabular MDP for planning, or for pre-training representations for later use in any deep RL based downstream task.
Significance of our work: AC-State uses simple multi-step inverse model along with bottleneck on the representations to learn latent states, which has not been exploited by prior works using dynamics models. Unlike the aforementioned works, here we focus on the representation learning problem. AC-State predicts actions based on future observations, and we show that a multi-step inverse dynamics model predicting actions can fully recover the control-endogenous latent states with no dependence on the exogenous noise. We emphasize that in the presence of exogenous noisc, methods based on predictive future observations are prone predicting both the control-endogenous and exogenous parts of the state space and do not have guarantees on recovering the latent structure. AC-State is the first practical and theoretically grounded approach that learns the control-endogemous representation with complex function approximators such as deep neural networks.
3 Exogenous Block MDP Setting
We consider the Exogenous Block Markov Decision Process (Ex-BMDP) setting to model systems with control-endogenous and exogenous subsystems, as formulated in Efroni et al. (202cc). This definition as- sumes that exogenous noise has no interaction with the control-endogenott state, which is stricter than the exogenous MDP definition introduced by Dietterich et al. (2018).
5

---

Under review as submission to TMLR
Notation and Assumptions: A BMDP consists of set of observations, X; set of latent states, Z; finite set of actions, A with cardinality A; transition function, T: 2 x A A(2); emission function 2 + A(X); a reward function R: : X' x - [0, 1]; and a start state distribution Ho E A(2). We do not consider the episodic setting, but only assume access to a single trajectory. The agent interacts with the environment, generating an observation and action sequence, (=1, #1, #1, R2, 22, 02, .2....) where =1 ~ ju(-). The latent dynamics follow 2+41 ~ T(x') 2t, 0+) and observations are generated from the latent state at the same time step: 2t ~ q(- 2x). The agent does not observe the latent states (2,,22)...), instead it receives only the observations (æ1, 222...). The block assumption holds if the support of the emission distributions of any two latent states are disjoint, supp(a(- =1) n supp(a(-|=2) = 0 when 21 # 22., where supp(a(- =)) = { E X g(x 2) > 0} for any latent state z. Lastly, the agent chooses actions using a policy n X' + A(A), so that Of ~ (-|2ae)-
We now define the model we consider in this work, which we refer as deterministic Ex-BMDP
Definition 3.1 (Deterministic Ex-BMDP). A determänistic Ea-BMDP is a BMDP such that the latent state can be decoupled into two parts z = (s, e) where E S is the control-endogenous state and E is the ezogenous state. We will further assume that the control-endogenous dynamics T(s' s, a) are deterministic. The ezogenous dynamicc Ta(e' |e) may be stochastic. For 2, x' E 2za E A the transition function is decoupled [ 22 =
T(51 18, s,a)Ta(e') (e).
The above dcfinition implics that there exist mappings f:: : the corresponding oodloo--ndoogenous and exogenous latent stales. The cardinality of the exogenous latent state E may be arbitrarily large. Furthermore, we assume that the diameter of the control-endogenous part of the state space is bounded. In other words, there is an optimal policy to reach any state from any other state in a finite number of steps:
- [S] and fr,e: - [E] from obscrvations to
Assumption 3.2 (Bounded Diameter of Control-Endogenous State Space). The length of the shortest path between any 21 E S to any 22 E S is bounded by D.
We now describe a structural result of the Ex-BMDP model, proved in Efroni et al.] (202cc)- We say that I is an endogenous policy if it is not a function of the exogenous noise. Formally, for any I1 and zg, if f. (211 ( (12) then T-- 21) = x(- (z2)).
Let P= s, be the probability to observe the control-endogennus latent state s f+('), t time steps after observing s' fa(r) and following policy æ, starting with action Let the exogenous state be defined similarly as e = (x) and c' = (z)) The following result shows that, when executing an endogenous policy, the future time step distribution of the observation process conditioning on any x has a decoupling property. Using this decoupling property we lator prove that the control-endogenous state partition is sufficient to minimize the loss of the AC-State objective.
Proposition 3.3 (Factorization Property under an Endogenous Policy, (Efroni et al. 2022c), Proposition 3). Assume that x ula) where j is some distribution over the observation space and that a is an endogenous policy. Then, for any 2 1 it holds that:
(x (2,,,)) = q(a' s'',e))P, [s, a, t)P((( e, t).
Additionally, we assume that the initial distribution at time step t = 0 is decoupled po(s, e) = (o(e)
J Justification of Assumptions: An intuitive justification for the concept of a control-endooenous latent state comes from the design of video games. A video game typically consists of a factorized game engine (which receives player input and updates a game state) and a graphics engine. Modern vidco game programs are often dozens of gigabytes, with nearly all of the space being used to store high-resolution textures, audio files, and other assets. A full generative model T(x++1 az) would need to use most of its capacity to model these high-resolution assets. On the other hand, the core gameplay engine necessary for characterizing the control-endogenous latent dynamics T(s+41 st. st)a)) takes up only a small fraction of the overall game program's 's size and thus may be much easier to mode To give a concrete example, when the 1993 game
Thhe control-endogenous latent state and their dynamics is generally even more compact than the game state, which itself may have exogenous noise such as unused or redundant variables.
6

---

Under review as submission to TMLR
Link's Awakening was remade in 2019 with similar gameplay but updated graphics, its file size increased from 8 MB to 5.8 GB, a 725x increase in size (Michel 2022). Successful reinforcement learning projects for modern video games such as AlphaStar (Starcraft 2) and OpenAl Five (Dota 2) take advantage of this gap by learning on top of internal game states rather than raw visual observations (Vinyals et al. 2019 Berner et al. 2019). In the real world, there no internal game state that we can extract. The control-endoeenee latent state must be learned from experience.
Additionally, we make the further assumption that the data is collected under endogenous policy. In this paper, we consider two kinds of policies: one is a random rollout policy, which is trivially an endogenous policy. The other type of policy we consider is achieved by planning to reach self-defined goals. This is achieved using a tabular-MDP constructed from counts of observed (s, 0, s') tuples, where 8 is the discretized output of our learned encoder. The endogenous policy assumption in our theory may not always hold in practice, as a learned encoder may have errors that cause it to depend on exogenous noise.
We have also assumed that the control-endogenous state space has finite state, bounded diameter, and deterministic dynamics. The assumption of finite state rules out continuous control problems and also makes problems with combinatorial structures (such multiple objects) infeasible. The requirement of a bounded diameter rules out environments that are unsafe or allow for irreversible changes, such as a robot, becoming permanently stuck or damaged. The assumption of deterministic oontrol-endogenote dynamics is not strictly required to use AC-State, but is required for a part of the proof that requires constructing the set of states that are reachable in a certain number of steps.
arm with many joints or robot interacting with
Goal: Our goal is reward-free and task-independent latent state discovery. Succossful discovery of the control-endogenous latent stale entails only learning a model for how to encode a such that T(8+11 | 8e, at), while not learning anything about T(e++1 ex) or how to encode e. In particular, we seek to show that the control-endogenous latent state is identifiable and can be recovered exactly in the asymptotic regime.
4 AC-State: Discovering Agent Control--Enogeenous Latent States
The control-endogennous latent state should preserve information about the interaction between states and actions while discarding irrelevant details. The proposed objective (Agent Control- Endirooo-n State or AC-
State) accomplishes this by generalizing one-step inverse dynamics (Pathak et al., 2017 Efroni et al. 2022c
Hutter& & Hansen 2022 Badia et al., 2020) to multiple steps with an information bottleneck (Figure) D This
multi-step inverse model predicts the first action taken to reach a control-endogenous latent state steps
the future:
Lac-state (F,z,a,t,)) k) -log(IP
(1)
The multi-step inverse objective Lac-state predicts the action a, from the observation just before the action a is taken and from an observation collected k steps in the future. The action step is sampled uniformly over all of the T steps in which the action was taken randomly. The prediction horizon k is sampled uniformly from 1 to a maximum horizon Ke, which may be a fixed hyperparameter or may depend on the number of steps left in the planning horizon at step t. The multi-step inverse model predicts the distribution over the first action to reach a state & steps in the future from a given current state. Generally, this action cannot be predicted perfectly, but still carries information about. the relationship between the current state and a state in the (potentially distant) future.
We optimize the parameters of an encoder model R" + {1...., Range(f)) that maps from an dimensional continuous observation to a finite latent state with output range of size Range(f), an integer number. FF is set that represents all mappings achievable by the model. We wish to minimize Objective while using an encoder with the smallest. latent stale as its output. Conceptually, we can define this as finding a set of optimal solutions. The endogeous latent state is guaranteed to be one of the optimal solutions for the multi-step inverse objective Lac-states while other solutions may fail to remove irrelevant information. The control-endogenons latent state uniquely achievable by finding the lowest capacity f in the set of optimal solutions to the AC-State objective:
7

---

Under review as submission to TMLR
G arg min E fer T- k))
E
Lic-scate (f, (t,
(2)
i E arg min Range SEG
(3)
As concrete strategy to achieve min- imal latent state, we combine two mech- anisms that are widely used in the deep learning literature to restrict the infor- mation capacity in f, which introduces an additional loss term Bottteeeck (f, z²). We first pass the hidden state at the end of the network through gaussian vari- ational information bottleneck (Alemi et al. 2017), which reduces the mutual information between a and the represen- tation. We then apply vector quantiza- tion (Van Den Oord et al., 2017). which yields a discrete latent state. While ei- ther of these could be used on their own, we found that adding the gaussian mu- tual information objective eased the dis- covery of parsimonious discrete repre- sentations. Both these are standard deep architectural elements, which we use along with their associated loss terms Wang 2020 Zeghidour al.) 2021). In-
Figure 2: AC-State Successfully Removes Exogenous Noise: In our experiments, we demonstrate an application of AC-State in a Robot Learning Task. AC-State discovers the control-endogenous latent state in a visual robotic sctting with temporally correlated distractors: TV, flashing lights, and drinking bird toys (left). The reconstruction from AC-State decoder captures the arm's position, while ignoring all the irrelevant background distractors (right). This latent representation captures the relevant parts of the world without any supervision, reward signal, or hard-coded assumptions.
Suitively, the bottleneck is necessary to ensure that the latent contains the minimal amount of information necessary to obtain the optimal solution to the multi-step inverse model. In the absence of a bottleneck, we could set the encoder to be the identity function Φ(x) = x and fail to learn a meaningful representation.
Since it avoids the need for complicated two-level optimization, we can consider using the bottleneck losses as additional terms in the objective. While we could re- re-weight these loss terms to make the bottleneck weaker, we did not find this to be necessary in practice. We optimize this approximate objective in practice:
f Re arg min FEJ beT k-U(IKK)
E
[CA--SEato U,I,,,tt, k) + Botttenecck xt) + Bottteeneck (f, setk)]
(4)
In addition optimizing the objective and restricting capacity, the actions taken by the agent are also important for the success of AC-State, in that they must achieve high coverage of the control-endogenous state space and not depend on the exogenous noise. This is satisfied by a random policy or policy that depends on f(zz) and achieves high coverage. The AC-State objective enjoys provable asymptotic success (see section 5 in discovering the control-endooeeous latent state.
Intuitively, the AC-State objective encourages the latent state to keep information about the long-term effect of actions, which requires storing all information about how the actions affect the world. At the same time, the AC-State objective never requires predicting information about the observations thersselves, so places value representing aspects of the world that are unrelated to the agent.
AC-State with Random Policy: The simpler version of AC-State uses a random policy to collect data and fit the encoder (Algorithmm D. While this is simple to implement, a uniform random policy is always exogenous-free free by construction and asymptotically achieves high coverage. At the same time, it may not achieve sufficient coverage with reasonable sample efficiency for some hard exploration tasks. If a single bad action can cause an agent to lose many steps of progress, the sample complexity needed achieve high coverage can scale exponentially in the horizon.
8

---

Under review as submission to TMLR
Algorithm 1 AC-State Algorithm for Latent State Discovery Using A Uniform Random Policy
Initialize observation trajectory x and action trajectory a. Initialize encoder Ja- Assune & control- endogenous diameter of K and number of samples to collect T, and set of actions A, and number of training iterations N.
2: 21 ~ U (μ(z))
3: for = 1, 2, ... T do
4: at ~ ~AA
5: Te41 ~ P(x [Tt, az)
6: for 1, ... N do
7:
1~ U~,,T) and k~ U1,
8:
== Lacsuuee (fa.t ()) + Bottleneck (fe. + Coottenncc (fo, Iz+k)
9: Update to minimize L by gradient descent.
AC-State with Planning Policy: A more involved algorithm (shown in Appendix C Algorithm 2 is able to discover latent states in problems where exploration is difficult (such as in an environment where specific actions need to be taken in sequence to make progress). This can be accomplished by using a planning policy instead of a purely random policy. While any type of policy could be used in principle, we experimented with planning in simple count-based tabular-MDP. In this model, monte-carlo breadth-firtt search (suitable for stochastic dynamics) is used to select rarely seen reachable states goal states. Then Dijkstra's algorithm is used to construct closed-loop plans to reach the desired goal state. All collected sample observations are collected into a buffer and their representations are recomputed under the current encoder network and the tabular- is reconstructed on every 100th iteration. Only a single random action is taken on the first stop of each goal-seeking trajectory, and it is this random action that predicted with the AC-State objective. While the AC-State theory assumes deterministic dynamics, the tabular-MDP constructed using learned latent states often has stochastic dynamics due to errors in the learned encoder.
We provide implementation details of the AC-State algorithm, with pseudo- code of the algorithm provided in C The AC-State algorithm collects data with either a uniform random rollout policy or a Dijkstra's 's algorithm based planning policy, as discussed in most of our experiments. This is because a random policy ensures that the actions are completely exogenous free, irrespective of whether the observation consists of exogenous information. We then learn an encoder representation using the multi-stop inverse dynamics objective as in equation D The encoded representation is then passed through Vector Quantization (VQ-VAE) bottleneck such to map the continuous representation into a set of discrete latent codes. The AC-State representation is learnt end to end using the overall objective, including the information bottleneck as in 4 We use the learnt discrete latents for visualizing the underlying structure of the MDP (Figure 0 right). The corresponding continuous representation is probed (using a separately learned decoder that predicts the ground truth state) to measure the accuracy of latent state recovery (figure D.. For each experiment, we include additional details data collection and the evaluation metric to measure accuracy of latent discovery.
Implementation Details:
Measuring Successful Discovery of Control-Endogennout Latent State: The result of training with AC-State is discrete graph where there is node for every control-endooeenous latent state with edges representing actions leading to a state-action node from which weighted outcome edges lead back to control- endogenous latent state nodes. We estimate the transition probabilities in this control-endogenous latent space by using counts to estimate T(s 4.). Once this control-endogenous latent state and associated transition distribution estimated, we can directly measure how correct and how parsimonious these dy- namics are whenever the ground-truth our experiments). We measure an Leerror on this dynamics distribution as well the ratio of the number of learned latent states to the number of ground truth control-endogenous states.
endooeonns latent state is available (which the case in all of
5 Theoretical Analysis
We present an asymptotic analysis of AC-State showing it recovers f., the control-endogenous state repre- sentation. The mathematical model we consider is the deterministic Ex-BMDP. There, the transition model
9

---

Under review as submission to TMLR
of the latent state decomposes into control-endogenous latent state, which evolves deterministically, along with noise term-the agent-irrelevnnt portion of the state. The noise term may be an arbitrary tempo- rally correlated stochastic process. If the reward does not depend on this noise, any optimal policy may be expressed in terms of this control-endogenons latent state. In this sense, the recovered control-endogenous latent stale is sufficient for achieving optimal behavior.
The proof involves steps. In the first step, we show that Bayes-optimal solution to the multi-step inverse model objective be achieved without dependence exogenous noise. In the second step, show that the minimal representation minimizing the multi-step inverse model objective is the control-endogenous state.
5.1 High Level Intuition for Theory
Our theorctical justification first proceeds by showing that the optimal multi-step inverse model only depends the control-endogenous state (Section 5.2 given our assumed factorization of the latent dynamics. Thus the control-endogenous latent state is sufficient for an optimal solution to the multi-step inverse model objective. In some sense, this establishes that the latent state does not need to be split any more than is necessary.
With h-1 there is no certificate to differentiate s1 and s2:
After moving from B by taking action "e",, sl is reachable in but s2 is not.
A
B
C
A
C
p(a-1 Ma, A;
s1
s2
p(a=J s1;
p(a-1 52, C)
s1
s2
Starting at B, this yields an h=2 certificate to differentiate and
p(a--1 A, -
p(a--1
A,
p(a-+1
B, C,C;
p(a-1 C,
pla-- B, p(a=- B, 2)
p(a=+1
B;
Figure 3: An illustration of the high-level idea of the proof in Section 5.3 The proof uses induction on h to produce certificates which differentiate all pairs of distinct states. An example of one certificate using 2-step inverse models is shown on the right side the figure. Going from B to requires the first action to move left (red arrow), whereas going to s2 requires that the first action moves right (blue arrow).
The next part of the argument (Section 5.3 and Figure 3) attacks from the opposite direction by showing that the multi-stop inverse model forces the latent state to split any two states which are truly distinct. This means that there is no coarser solution to the multi-step inverse model objective (since the bottleneck will prefer the coarsest solution). The proof uses induction on the horizon length to produce certificates split increasingly distantly-separated states. This novel result is somewhat counter-intuitive in a few ways. One surprising result is that it only requires the model to predict the first action taken reach future state, rather than predicting the full sequence of actions. Another interesting result is that the policy following the first action can be fairly general (needs to achieve high state coverage over the dataset) and includes a purely random rollout policy. The optimal solution to the multi-stop inverse model objective will chango depending on the choice of rollout policy, but the theory indicates that a wide variety choices will lead to learning the correct representation.
5.2 Bayes Optimal Solution does not depend on Exogenous Noise
Consider the generative process in which a is sampled from a distribution μ, the agent executes a policy a for time steps and samples z'. Denote by (x,, x',t) .t) as the joint probability, and by (x x, x,2',t) as the probability that under this generative process the action upon observing x a. The following result, which builds on Proposition 3.3 shows that the optimal bayes solution Pr,a (a x, x' is equal Pa,,(a fx for P fx x² 0, where P.
Pr (Ir, x',
is the probability to sample x
10

---

Under review as submission to TMLR
Proposition 5.1. Assume that is an endogenous policy. Let i for some distribution p. Then, the Bayes' optimal prediclor of the aclion-prediction model is piece-wise conslant with respect to the control- endogenous partition: for all a E A, t 0 and I, x' E & such that IP P
(I, x', b) > 0 it holds that
(a x 2, .)) (a (.()))
We comment that the condition Pr. (2, x', tt) 0 is necessary since otherwise the conditional probability P (a Ix,
is not well defined.
The proof for this proposition is a Miraightforward application of bayes theorem, is given in the appendix, and was previously presented in Efroni et al. (002cc a).
5.3 Discovery of Full Control-Endogenonn Latent State
Proposition 5.1 from the previous section shows that the multi-step action-prediction model is piecewise constant with respect to the partition induced by the control-endogenous states f. X + [S]. In this section, we assume that the exccuted policy an endogenous policy that induces sufficient exploration. With this, we prove that there is no coarser partition of the observation space such that the set of inverse models is piecewise constant with respect
We assume that the Markov chain induced on the control-cndogenont state space by exccuting the policy ap by which AC-State collects the data has stationary distribution μD such that pp(s,a) > 0 and TD(a s) 2 Wmin for all E S and E A. We consider the stochastic process in which an observation is sampled from a distribution such that μ(s) (s(ss for all s. Then, the agent executes the policy TD for t time steps. For brevity, we denote the probability measure induced by this process as Pp. A more detailed justification for this assumption is given in the appendix.
We begin by defining several useful notions. We denote the set of reachable control-endogenous states from in h time steps as (s).
Definition 5.2. (Reachable Control-endogenoos States). Let the set of reachable control-endogenous states from 8 in h > 0 time steps be RRss, h) {s' max, (s h,h)=11
Observe that every reachable state from s in time steps satisfies that max, F (s' | = s, due the deterministic assumption of the control-endooennus dynamics.
Next, we define notion of consistent partition with respect to set of function values. Intuitively, a partition of space X' is consistent with a set of function values if the function is pieco-wisc constant on that partition.
Definition 5.3. (Consistent Partition with respect to g). Consider set G (o(a, v. y')) wo ae where A D x y + [0, We say that f y + [N] is consistent partition with respect to G if for all 3. J/2 E ,, f(y)) = f(1)) implies that g(a, v.3() = g(a, y, v2) for all a E A.
Observe that Proposition 5.1 shows that the partition of X' according to f. is consistent with respect to (P x, h) x. E Ex [H] s.t. Po(x,2', > 0}, since, by Proposition 5.1 Pp | x,,', h) = (P
(x'),,
Towards establishing that the coarsest abstraction according to the AC-State objective f. we make the following definition.
Definition 5.4. (The Generalized Inverse Dynamics Set AC(s, h)). Let S, E N. We denote by AC(s, the set of malli-step inverse models accessible from in h time sleps. Formally,
AC(s, h) (Pp(a) s', s" h') : s' R(s, s" e
A, [h]]
Observe that equation equation the inverse function Po(a ,,,"",h') is always well defined since Pp(s', s", ,h) > O. It holds that P s.h"" = Pp(s')Ppos" s', s/,h') > 0, since Po(s') > and Pp(s" s², > 0. The inequality PD(s')= > 0 holds by the assumption that the stationary distribution when following
11

---

Under review as submission to TMLR
U has positive support all control-endogenone. states (Appendix D.3 The inequality PD(s" s' > 0 holds since, by definition s" E (R,s, (() is reachable from s' in h time steps; hence, PD(s" 1 > 0 by the fact that the stationary assumption (Appendix D.3) implies that the policy TD induces sufficient exploration.
Theorem 5.5. (f+ is the coarsest partition consistent with respect to AC-State objective). Given the bounded diameler Assumplion 3.2 and the high-coverage policy assumption (Appendiæ D.3) hold, il follows that there is no coorser partition than Sa, which is consistent with AC(s, D) for any 8 E ES.
The proof of this theorem, which uses induction on h, is provided in the appendix.
6 Experiments
Summary of Experimental Results: Our experiments explore three domains and demonstrate the unique capabilities of AC-State: (i) AC-State learns the control-endogenous latent state of real robot. from high-resolution video of the robot with rich temporal background structure (a TV playing a video, flashing lights, dipping birds, and even people) occluding the information in the observations; (ii) AC-State learns about the controlled agent while ignoring other random agents in a multiple maze environment where there are other functionally identical agents. We demonstrate that AC-State solves a hard maze exploration problem. (iii) AC-State learns control-endogenous latent state in house navigation environment where the observations are high-resolution images and the camera's vertical position randomly oscillates, showing that AC-State is invariant to exogenous viewpoint noise that radically changes the observation.
6.1 Robotic Arm Experiment
We first evaluate AC-State to recover the latent states in real robot arm experiment consisting of ob- servations with exogenous information in the background. We briefly describe the experiment setup and the evaluation metric for measuring latent state recovery. We found that AC-State discovers the control- endogenous latent state of real robot arm while ignoring distractions.
Agent's view
Top vlew
t=1
t=2
t=3
t=2
t=3
80
70
60
50
Autbencoder
Error
40
Inverse
AC-State
30
20
10
1
14
27
40
53
66
72
92
Training epochs
Figure 4: A real robotic arm moves between nine different positions (left). The four actions correspond to the arm moving forward, backward, to the left, or to the right. The quality of control-cndogenous latent state dynamics learned by AC-State is better than one-step inverse models and autocncoders, as measured by the dynamics difference crror (bottom right).
We collected 6 hours of data (14,000 samples) from an AR-3 robot arm (Annin 2022) by taking give high-level actions (move left, move right, move up, move down and stay). A picture of the robot was taken after each completed action. We used maximum horizon of corresponding to the diameter of the endogenous part of MDP, but empirically observed that larger horizons also worked
Experiment Setup:
12

---

Under review as submission to TMLR
well. The observations were mapped to representation by using 6-layer vision transformer (Dsovvisskiy et al. 2021) with 4-heads per-layer. The action prediction for AC-State used a simple MLP on top of these representations and an embedding for k. Details of our experiment setup are provided in Appendix A.3 The robot was placed among many distractions, such as a television, flashing, color-changing lights, and people moving the background. Some of these distractions (especially the TV) had strong temporal correlations between adjacent time steps, as often the case real-life situations.
shows that AC-State discover the control-endogemous latent state which is the ground truth position of the robot (not used during training). Qualitatively, we measure the accuracy of latent recovery of AC-State by training convolutional neural network to reconstruct x from f(x) with square loss as the error metric. In other words, we learn & visualization of the latent state to map from the latent state f(z)) to an estimate ât the observation 2; by optimizing the mean-square error reconstruction loss | - -0-=x,1' We found that the robot arm's position was correctly reconstructed, while the distracting TV and color-changing lights appeared completely blank, as expected (Figure 2). As discussod previously, auto-encoder trained end-to-end with as input captures both the control-endooenouu latent state and distractor noise. The theoretical counter-examples for autoencoders and one-step inverse models also closely match the crrors that we observe experimentally (Figure Our results show that AC-State can perfectly recover control-endogenout latent states, both qualitatively and quantitatively, in a practical robot arm task consisting of exogenous information in the background.
Experiment Result and Evaluation Figure S
180
304
3.1
03
(10)
0.8
200
10
1220
(V0
0.8
233
102
1000
RUD
211
(a) Autoencoder (Theory worst-case)
(b) Inverse
(c) AC-State
(Theory worst-case)
(Theory worst-case)
0.9
X0
09
83
0.0
201
0.8
03)
0.6
00
2.1
R2
(d) Autoencoder
(e) Inverse (Empirical)
(f) AC-State
(Empirical)
(Empirical)
Figure 5: In co-occurrence histograms measuring performance the robotic arm environment, autocncoders fail (left), one-step inverse models fall prey to the same counterexample in theory and experiment (center), and AC-State discovers a perfect endroo-endogenous latent state (diagonal histogram, right). Intriguingly, the failure case of the one-step inverse model is the merging of three distant states, which are the arm being the center-left, center, or center-right positions. These states being merged by one-step inverse model is precisely what the theory predicts. On the other hand, the auloencoder merges states where the robot arm is in similar positions because the observations are visually similar.
6.2
Mazes with Exogenous Agents and Reset Actions
In our experiments, AC-State removes exogenous noise in environments where that noise has a rich and complex structure. We studied a multi-agent system in which a single agent is control-endogenous and the other agents follow their own independent policies. To do this, we take pixel-basod visual gridworld consisting of multiple mazes and agents, where the goal of AC-State is to recover only the controllable part of the maze while ignoring observations from other agents. Tn our experiments, we again qualitatively
13

---

Under review as submission to TMLR
show that AC-State perfectly recover the underlying structure of the endogenous, controllable maze. Additional quantitative results measuring accuracy of latent state recovery is provided in appendix A.2.
Experiment Setup: In an environ- ment with agents and each agent hav- ing e control-endogenous states, the over- all size of the observation space ca. With 3,000 training samples, AC-State is able to nearly perfectly discover the agent's control-cndogenous latent state, while fully ignoring the state of the un- controllable exogenous agents, with all of the agents controlled by a random policy (Figure When using a random pol- iey, the maximum horizon K is set to the diameter of the MDP, whereas when us- ing AC-State the maximum horizon is based on the number of steps to the goal estimated by the planning process. Addi- tional details on the multi-maze environ- ment is provided in appendix A.2
1000
2000
3000
4000
5000
6000
10000
15000
20000
25000
1000
2000
3000
4000
5000
6000
10000
15000
20000
25000
Experiment Result and Evaluation: In the multi-maze system consisting of exogenous observation from other agents, we first evaluate the ability of AC-State to recover the underlying structuro the controllable states. Figure G shows that using the AC-State objective followed by the discretization bottleneck, sualize the endogenous states with set of discrete latent codes using the quantization. Figure 6 (top-left) shows the observation of the agent with the controllable maze (top-middle) highlighted. Figure O (top-right) shows the extracted visualization with the discrete codes, qualitatively showing that endogenous-statee can be recovered to retrieve the structure the controllable MDP. Our goal to show that the proposed AC-State agent can "discover" the control-endogenous latent state within the global gridworld while ignor- ing the structural perturbations the geometry of the other gridworlds. comparison with basclines is shown in Table 2
Figure 6: We study multi-agent world where each of the nine mazes has separately controlled agent. Training AC-State with the actions of specific agent discovers coutrol-endogenous latent information about the other mazes (top). In version of this environment where a fixed third top-left corner, random policy fails explore (top hcatmap sequence), whereas planned exploration using AC-State reaches all parts the maze (bottom heatmap sequence).
while discarding
cause the agent's position to reset to
vi-
Algorithms
TRUE ENDOGRNOUS STATES EXPLORED
Dynamics Difference Frror (%)
Contrastive
20/68
92.2
Autoencoder
20/68
91.7
Forward Generative
25/68
97.1
AC-State
68/ /6
0.00
Table 2: Multiple-Matt Maze Experiment Result with Baseline Comparisons, with each model having 25000 samples to explore in the environment. The baseline methods fail to ignore exogenous noise, fail to learn an offective tabular representation, and thus fail explore effectively.
We then show that the control-endogenont latent state is useful when it allows for exponentially more efficient exploration than is achievable under a random policy. To exhibit this, we modified the maze problem by giving the agent additional actions that reset. to a fixed initial starting position. When a third of all actions cause resets, the probability acting randomly for N steps without resetting (2/3) We show that a lcarned exploration policy using AC-State succoods in full exploration and learning of the control-cndogenous latent state with 25,000 samples, while random exploration policy barely explores the first room with the number of samples (Figure
14

---

Under review as submission to TMLR
The counts of the discrete latent states are used to construct simple tabular MDP where planning is done to reach goal states using monte carlo version of Dijkstra's algorithm (to account for stochastic transition dynamics). The reachable goal states are sampled proportionally to count(s.]. so the rarely seen states are the most likely to be selected as goals. Experiment results demonstrate that a goal-seeking policy achieves perfect coverage of the state space by using discovered latents for exploration, while a random policy fails to reach more than 25% of the state space in the presence of reset actions. We demonstrate this with heatmaps showing state visitation frequencies.
6.3 First-Person Perspective House Navigation
Finally, we evaluate AC-State in a simulated houschold navigation task where data is pre- collected from different, positions in the house- hold. We compare AC-State with several base- lines and evaluate the ability to capture the endogenous latent state, ie, position of the agent the house while discarding other irrel- evant information captured by the high resolu- tion camera. Additional details on the matter- port simulator, along with experiment setup and implementation details, are provided in appendix A.4
Experiment Setup: : In order to analyze the performance of the proposed AC-State objec- live in more realistic setting, we evaluated it Matterport (Chang et al. 2017). navi-
gation environment where each observation is
high resolution image taken from position in a real house. A 20,000 sample dataset col- lected from an agent moving randomly through the house used to train AC-State. In addi- tion to the high degree of visual information in the input observations, we randomly move the camera up down at each step as controlled source of irrelevant information (exogenous noise). AC-State removes the view noise from the encoded reprosentation f(x while still cap- turing the true control-enndogenous latent state (Figure 77. whereas other baselines capture both the control-endogenous latent state and the exogenous noise. The first- fiss-person house images are mapped to representations using a 6-layer, 4-head vision transformer (Dosovitskiy et al. 2021). The model is trained end-to-end with the AC-State objective with maximum horizon of K for 20 epochs using the Adam optimizer with a learning rate of 1c-4. We ran three soods and report mean and standard deviation. Details the matterport simulator are provided in Appendix A.4
State parsimony
Accuracy
Contro lable
Esagerous aebeeggponng socurey
100
Score
Adto Contrestive
Gomeratten mosill
wid anttececc
AC-Budo
Cemerativo
wi butteracc
Aisdo
AC vists
Auto
Figure 7: We evaluate AC-State in house navigation environment (top), where the agent observes high-resolution images of first-person views and the vertical position of the camera exogenous (center). AC-State discovers a control-endogenous latent state that is parsimo- nious (bottom left)) and captures the position of the agent in the house while discarding exogenous noise (bottom right). The basclines capture the control-endogenoos latent state but fail discard the ex- ogenous noise.
Experiment Results and Evaluation We present the results for this experiment in Figure 7 The Con- trollable Latent State Accuracy is the viewpoint prediction accuracy for the current state. The Ezogenous Noise-Ignoring Accuracy reflects how much information about the exogenous noise removed from the ob- servation. We can see that the proposed AC-State model has the highest control-endogenous latent state and exogenous noiso-ignoring accuracy. Thus, it outperforms the basclines we considered at capturing control- endogenous state while ignoring exogenous noise. We calculated state parsimony as Num. Ground Truth States Num. Discovered States Therefore, a lower state parsimony denotes high number of discovered states, which means that the model
15

---

Under review as submission to TMLR
fails at ignoring exogenous information. The proposed model has the highest state parsimony, which shows the effectiveness of the model in ignoring exogenous noise while only capturing control-endogenous latent states.
7 Conclusion
AC-State reliably discovers the control-endogenous latent state across multiple domains. The vast simplifi- cation of the control-endogenous latent state discovered by AC-State enables visualization, exact planning, and fast exploration. The field of self-supervised reinforcement learning particularly benefits from these approaches, with AC-State useful across wide range of applications involving interactive agents as self contained module to improve sample efficiency given any task specification. As the richness of sonsors and the u computing technologies (such as virtual reality, the internet of things, and self-driving cars) continues to grow, the capacity to discover agent-control endogenous latent states enables new classes of applications.
Limitations and Future Work: In this paper, we showed that AC-State recover the control en- dogenous latent states with perfect accuracy, and recover the underlying latent tabular MDP by completely discarding the exogenous information. However, one limitation is that the success of AC-State depends on the use of the discrete information bottleneck and only applicable for environments with deterministic endogenous dynamics. An interesting extension of AC-State would derive an algorithm, with provable guarantces for stochastic environments and starting states. Furthermore, would be interesting to see how the learnt representations from AC-State can be useful for the dewnstream RL task in presence of exogenous information. Since AC-State recovers the full underlying model, it can potentially be used for reaching any goal state in the exploration frontier, making the algorithm highly applicable for practical tasks such as household navigation.
16

---

Under review as submission to TMLR
References
Alexander A Alemi, lan Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottle- neck. Proceedings of the International Conference on Learning Representations (ICLR), 2017.
Chris Annin. Ar3 robot, 2022. Location: hutps//wwww.annnnrbbiis.oomm
Adríà Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Zhao- han Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. In Inter- national Conference on Machine Learning, pp. 507-517. PMLR, 2020.
Bowen Baker, Ilge Akkaya, Peter Zhokhov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. arXiv preprint arXim:200. 11795, 2022.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemystaw Dębiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv: 1912.
1912.06600, 06680, 2019.
Homanga Bharadhwajj, Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. Information pri- oritization through empowerment in visual mosel-based RL. In International Conference on Learning Representations, 2022. URL htttps://ppenrevier http://tw....pppeeeeeeeeee/////a http:://oormmmm/dd.b//yyyao1
Cesar Cadena, Luca Carlone, Henry Carrillo, Yasir Latif, Davide Scaramuzza, José Neira, Ian Reid, and John J Leonard. Past,, present, and future of simultancous localization and mapping: Toward the robust- perception age. IEEE Transactions on Robotics, 32(6):1309-1333 2016.
Angel Chang, Angela Dai, Thomas Funkhouser, Maciej Halber, Matthias Niessner, Manolis Savva, Shuran Song, Andy Zeng, and Yinda Zhang. Matterport3D: Learning from RGB-D data in indoor environments. International Conference on 3D Vision (3DV), 2017.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PML.R, 2020.
Kingma Diederík, Ba Jimmy, et al. Adam: Method for Stochastic Optimization. International Conference on Learning Representations, pp. 273-297, 2014.
Thomas G. Dietterich, George Trimponias, and Zhitang Chen. Discovering and removing exogenous state variables and rewards for reinforcement learning. CoRR, abs/1806.0135 01584, 2018. URL htRL::t////axxvoorr abs/1806. 01584
MWM Gamini Dissanayake, Paul Newman, Steve Clark, Hugh F Durrant- Durrant-Whyte, and Michael Csorba. A solution to the simultancous localization and map building (slam) problem. JEEE Transactions on Robotics and Automation, 17(3):229 241, 2001.
Alexcy Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkorcit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations, 2021.
Hugh Durrant-Whyte and Tim Bailey. Simultaneous localization and mapping: part i. IEEE Robotics B Automation Mogazine, 13(2):99 110, 2006.
Yonathan Efroni, Dylan J Foster, Dipendra Misra, Akshay Krishnamurthy, and John Langford. Sample- efficient reinforcement learning in the presence of exogenous information. (Accepted for publication at) Conference on Leurning Theory, 2022a.
Yonathan Efroni, Sham Kakade, Akshay Krishnamurthy, and Cyril Zhang. Sparsity in partially controllable linear systems. (Accepted for publication at) International Conference Machine Learning, 2022b.
17

---

hwtp://wwdd.....or.r..eeeeeeeeee Under review as submission to TMLR
Yonathan Efroni, Dipendra Misra, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Provably filtering exogenous distractors using multistep inverse dynamics. In International Conference on Learning Representations, 2022c. URL https://openreviiet http:///wwee.tttoooooorddddedeee
Xiang Fu, Ge Yang, Pulkit. Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 3490- PMLR, 18-24 Jul 2021. URL
httts::/proceeding mlr. www..reesssvvvii press/v139/fulib html
lan Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep Learning, volume 1. MIT Press, 2016.
Zhaohan Daniel Guo, Shantanu Thakoor, Miruna Pislar, Bernardo Avila Pires, Florent Altché, Corentin Tal- lec, Alaa Saade, Daniele Calandriello, Tean-Bastionn Grill, Yunhao Tang, et al. Byol-explore: Exploration by bootstrapped prediction. arXiv preprint arXiv:2206. 08332, 2022.
Dan Hendrycks and Kevin Gimpel. Gaussian Error Linear Units (GELUs), 2016.
Marcus Hutter and Steven Hansen. Uniqueness and complexity of inverse mdp models, 2022.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameteerization with gumbel-softmax. In 5th Interna- tional Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.nel, 2017. URL
Leslie Pack Kaelbling, Michael L. Littman, and Anthony R. Cassandra. Planning and acting in partially observable stochastic domains. Artif. Intell, 101 (1-2):99-134 1998. doi: 10.1016/S0004-3700 03000443700299)000033
URL http:htttsss..orgggggee/ http://///......o/ooooo/////0///
10t.111111100000 (9)))000000000))))))0) (98) 00023-X
Alexander S. Klyubin, Daniel Polani, and Chrystopher L. Nehaniv. All else being equal be empowered. Mathieu S. Capcarrère, Alex A. Freitas, Peter J. Bentley, Colin G. Johnson, and Jon Timmis (eds.).
Advances in Artificial Life, pp. 744 753, Berlin, Heidelberg, 2005. Springer Berlin Heidelberg. ISBN 370-3-440-31
Sascha Lange, Martin A. Riedmiller, and Arne Voigtländer. Autonomous reinforcement learning on raw visual input data in a real world application. In The 2012 International Joint Conference on Neural Networks (IJCNN), Brisbane, Australia, June 10-15, 2012, pp. 1-8. LEEE, 2012. doi: (00p:///JCNNNN012.665223325 URL https://doi. org/10. 1109 2012. 6252823
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. Nature, 521((553)::366 444, 2015.
David A Levin and Yuval Peres. Markou chains and mixing limes, volune 107. American Mathematical
Soc., 2017.
Yunzhu Li, Antonio Torralba, Anima Anandkunarr Dieter Fox, and Animesh Garg. Causal dis- covery in physical systems from videos. In Н. Larochelle, M. Ranzato, R. Hadsell, M.F. Bal- can, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 9180 9192. Curran Associates, Inc., 2020. URL 6822951732be44eee htt.aeeeeeddeaaaaaaaaaaaaaaaaaaa
pdf..
Michael L. Littman, Anthony R. Cassandra, and Leslie Pack Kaelbling. Learning policies for partially observable environments: Scaling up. In Armand Prieditis and Stuart Russell (eds.), Machine Learning, Proceedings of the Twelfth International Conference on Machine Learnång, Tahoe City, California, USA, July 9-12, 1995, pp. 362 370. Morgan Kaufmann, 1995. doi:
10000//888--555-------0000-22200 URL
sttpps///aos
Michael L. Littman, Richard S. Sutton, and Satinder Singh. Predictive representations of state. In Thomas G. Dietterich, Suzanna Becker, and Zoubin Ghahramani (eds.), Advances in Neural Information Process- ing Systems 14 (Neural Information Processing Systems: Natural and Synthetic, NIPS 2001 December 3-8, 2001, Vancouser, British Columbia, Canadaj, pp. 1555-1551111 MIT Press, 2001. URL https:// proceedings http:///wwsaaaa.//oo////////aaaa hht
18

---

Under review as submission to TMLR
Lennart Ljung. System Identification. In Signal Analysis and Prediction, pp. 163-173. Springer, 1998.
Atalanti A Mastakouri, Bernhard Schölkopf, and Dominik Janzing, Necessary and sufficient conditions for causal feature selection in time series with latent common causes. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 7502-7511. PMLR, 18-24 Jul 2021. URL htttps://proceediggg mlr.
Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and R. Devon Hjjelm Deep reinforcement and infomax learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Had- sell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Newral Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NewrIPS 2020, December 6-12, 2020, virtual, 2020. URL https://procceedings wewprdiinnnnnniunispe neurips www.apps...ppppre/22222222..22 html
Teddy Michel The file size of every core legend of zelda game, 2022. URL https://www. wwwpwwwwwwwwwww.... comm
Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state abstraction and provably efficient rich-observation reinforcement learning. In International conference on machine learning, pp. 6961 6971. PMLR, 2020.
Volodyuyyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglon, Daan Wienstra, and Martin Riedmiller. Playing Atari with Deep Reinforcement Learning. Neural Information Processing (NIPS).
Shakir Mohamed and Danilo Jimenez Rezende. Variational information maximisation for intrinsically motivated reinforcement learning. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 28: An- nual Conference Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pp. 2125 2133, 2015. URL https://proceeeding/ wetpeeediinnnnuuuis....o /ttp:/ppp..pppppp/////////////// 404404444eeeee55574afff77f644444
Abstract. html
Pierre Yver Oudeyer and Frédéric Kaplan. What is intrinsic motivation? A typology of computational approaches. Frontiers Newrorobotccs, 1:6, 2007. doi: 12t.9//nnrro...0..220000
12.006.2007. URL, https://oo/:ooor/
10. 10.3389/neuro 12.006.222 2007
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity- exploration by self- supervised prediction. In International conference on machine learning, pp. 2778-2788 PMLR, 2017.
Kate Rakelly, Abhishek Gupta, Carlos Florensa, and Sergey Levine. Which mutualliinfrmation representa- tion learning objectives are sufficient for control? In Advances in Newral Information Processing Systems, volume 34, pp. 26345-263311 Curran Associates, Inc., 2021. URL https://proceeddinn htwp:/eeddding....uis...... neurips. cc/ [ppp://0001ttttadd666eeeeeeeeee2
Nino Scherrer, Olexa Bilaniuk, Yashas Annadani, Anirudh Goyal, Patrick Schwab, Bernhard Schölkopl, Michael C Mozer, Yoshua Bengio, Stefan Bauer, and Nan Rosemary Ke. Learning neural causal models with active interventions. arXiv preprint Xrxiv:2109. 02.0020, 2021.
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bachman. Data-efficient reinforcement learning with peed-predictive representations. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 8-7, 2021. OpenReview.net, 2021a URL
Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R. De- von Hjelm, Philip Bachman, and Aaron C. Courville. Pretraining representations for data-efficient re- inforcement learning. In Marc'Aurcli Ranzato, Alina Beygekzmer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: An- nual Conference on Neural Information Processing Systems 2021, NewrIPS 2021, December 6-14,
19

---

Under review as submission to TMLR
2021, virtual, pp. 12686- 2021b. URL http:://ppoeeiig....rrrp...opppp 69eba3467163effeee beeaaaaa 3eeaeeeeaaaaaaaaa1aaa00a00
html
Lin Song, Peter Langfelder, and Steve Horvath. Comparison of co-expression measures: mutual information,
correlation, and model based indices. BMC Bioinform., 13:328, 2012. doi: 10.1186/1471-21055 21171--210--1-3-222 13-328.
URL https://dol org/10.
Bradly C. Stadie, Sergey Levine, and Pieter Abbeel. Incentivizing exploration in reinforcement learning with deep predictive models. CoRR, abs/ abs/1507.00814 2015. URL
Jithendaraa Subramanian, Yashas Annadani, Ivaxi Sheth, Nan Rosemary Ke, Tristan Deleu, Stefan Bauer, Derek Nowrouzezahrai, and Samira Ebrahimi Kahou. Learning latent structural causal models. arXiv
preprint arXi::2210. 13583, 2022.
Ilya O. Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucie, and Alexey Dosovitskiy. Mip-mixer: An all-mlp architecture for vision. Advances in Neural Information Processing Systems, NewrIPS 2021, December 6-14, 2021, virtual, pp. 24261-24272, 2021. URL https://pproeeddin aiatfee-bbbrtatt.ht.a
Aaron van den Oord, Oriol Vinyals, and koray kavukcuoglu. Neural discrete representation learning, In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnelt (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https:
7/proceedings wewpeddiinnnnnuuiiss..
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedinng het.:eseerrrisccccccppppo hearrips.c//ppeer221///11 31be224647/6e1111053.cccca.6.. Pappee..
Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782):350 350-354, 2019.
Niklas Wahlström, Thomas B. Schön, and Mare Peter Deisenroth. From pixels to torques: Policy learning with deep dynamical models. CoRR, abs//150.00225 2015. URL http:p///aarii..rr//////0.......
Phil Wang. Vector quantization pytorch. http:////tthhccooooccccc//////// 2020.
Tongzhou Wang, Simon S Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised mdps: Learning world models better than the world itself. arXiv preprint arXiv:2206.15// 2022a.
Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for task- independent state abstraction, 2022b. URL https://arxiv.. wwwp:/s.g//aa..oe.. htsp/bbb/2206.. 13452
Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin A. Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pp. 2746 2754, 2015. URL Awte--atsttttttteeommmmoom.. html
neuripg............e.. cc/
paper
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 34:25476 24:25476-25488, 2021.
Tianhe Yu, Gleb Shevchuk, Dorsa Sadigh, and Chelsea Finn. Unsupervised visuomotor control through distributional planning networks. Robotics: Science and Systems (RSS), 2019.
20

---

Under review as submission to TMLR
Neil Zeghidour, Alejandro Luebs, Ahmed Omran, Jan Skoglund, and Marco Tagliasacchi. Soundstream: An end-to-end neural audio codec, 2021.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.nec, net, 2021a. URL http::///pppppees...eeee//re///a
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations ions for reinforcement learning without reconstruction. In International Conference Learning Representations, 2021b. URL http:s///ppaaaerreee...e/rrrrrro 2FCWDKRREU
21

---

Under review as submission to TMLR
A Methods and Experiments
A.1 Dynamics Difference Error Metric
For every pair of learned state and action (s, a), we have an empirical count-based distribution over next learned states P(s'|s, a). We can also empirically estimate the distribution over ground states for each learned state, which translated from the learned state dynamics. We can also look at the empirical transition distribution over ground states. We then compute an L1-difference between these two distributions and uniformly average over all pairs of ground state and action values and refer to this as the dynamics difference error.
call P(g]s). We can sample over this to get an expected dynamics over the ground states
In this section, we will describe our methods and experiments for validating the proposed latent state discovery with AC-State. Three experiments, including both simulated and physical environments, are used test the efficacy of our proposed algorithm. The environments are carefully chosen to demonstrate the ability of the AC-State agent to succeed at navigation and virtual manipulation tasks with varying degrees of difficulty; on these lestbeds, algorithuns with similar properties in literature fail to succeed. In what follows, we will describe the environmenall setups, the function approximation scheme for the latent state, and the results that we produced.
A.2 Mazes with Exogenous Agents and Reset Actions
We consider global 2D maze (see Fig- 6 further divided into nine 2D maze substructures (henceforth called gridworlds). Each gridworld is made up of 6 x 6 ground truth states, and only one of the gridworlds contains the AC-State agent. Every gridworld other than the one containing the true agent has an agent placed within it whose motion is governed by random actions. An ablation where we vary the maximum number of learned codes available to the model is shown in Table 3 A table showing progress in exploration (along with the number of learned codes being used) is shown in Table
Number of Codes
Ground Truth States
% STATES EXPLORED
% DYNAMICS DIFFERENCE FRROR
50
68
80.9
25.4
60
68
97.1
8.7
68
98.5
4.4
80
68
100
0.0
Table 3: Varying number of codes (learned states) available to the model while using AC-State on the 4-room multiple maze task. The percentage of ground states explored after 25000 samples and how the dynamics error varies as we increase the number of codes. AC-State typically needs more states than are actually present to fully succeed, yet performance degrades gracefully if the number of learned states is too restricted.
A.2.1 Exploring in Presence of Reset Actions
Data Collection: We collect data under a random roll-out policy while interacting with the gridworld's environment. We endow the agent with the ability to "reset" its action to a fixed starting state. The goal of this experiment is to show that, in the presence of reset actions, it is sufficiently hard for a random rollout policy to get full coverage of the mazes. To achieve sufficient coverage, control-endogenots latent states to learn goal seeking policy that can be incentivized to deterministically reach unseen regions of the state space. The counts of the discrete latent states are used to construct simple tabular MDP where planning is done to reach goal states using monle carlo version of Dijkstra's 's algorithm (to account for stochastic transition dynamics). Experimental results demonstrate that a goal-seeking policy achieves perfect coverage of the state space by using discovered latents for exploration, while a random policy fails to reach more than 25% of the state space in the presence of reset actions. We demonstrate this with heatmaps showing state visitation frequencies.
can leverage the discovered
22

---

Under review as submission to TMLR
Number of Samples
GROunD TRUTH STATES EXPLORED
NUMBER LEARNED STATES USED
2000
27
36
4000
47
57
6000
8000
63
70
10000
63
72
12000
66
75
14000
77
16000
66
77
18000
68
79
20000
77
Table 4: Using the AC-State algorithm on the 4-room maze exploration task, we show how the number of seen ground truth states and the number of learned codes change, as we increase the number of samples seen during the exploration process.
Experiment Details: The encoder receives observations of size 80 x 720 x 3 due to the observations from other exogenous agents. The agent has an action space of 4, where actions are picked randomly from uniform policy. For the reset action sctting, we use an additional reset actions, and uniformly picking a reset action can reset it to a deterministic starting state. The observation is encoded using the MLP-Mixer architecture (Tossikkiin et al.) 2021 with gated residual connections (Jang et al., 2017p. The model is trained using the Adam optimizer (Diederik et al. 2014) with a default learning rate of 0.0001 and without weight decay. We use a 2-layer feed-forward network (FFN) with 512 hidden units for the encoder network, followed by a vector quantization (VQ-VAE) bottleneck. The use of a VQ-VAF bottleneck would discretize the representation from the multi-step inverse model by adding a codebook of discrete learnable codes. For recovering -ndogeonns latents from the maze we want to control while ignoring the other exogenous mazes, further use MLP-Mixer architecture Tolstikhin et al. (2021) with gated residual updates Jang et al. (2017). Both the inverse mode and the VQ- VAE bottlenecks updated using an Adam optimizer Diederik et al. 20114) with default learning rate of 0.0001 without weight decay.
A.3 Robotic Arm under Exogenous Observations
Using robotic arm (Annin 2022) with 6 degrees of freedom, there are 5 possible abstract actions: forward, reverse, left, right, and stay. The robot arm moves within 9 possible positions in a virtual 3x3 grid, with walls between some cells. The center of each cell is equidistant. from adjoining cells. The end effector is kept at constant height. We compute each cell's centroid and compose a transformation from the joint space of the robot to particular grid cells via standard inverse kinematics caleulations. Two cameras are used to take still images. One camera is facing the front of the robot, and the other camera is facing down from above the robot. When a command is received, the robot moves from one cell center to another cell center, assuming no wall is present. After each movement, still images (640x480) are taken from two cameras and appended together into one image (1280x480). During training, only the forward facing, down-sampled (256x256) image is used. Fach movement takes one second. After every 500 joint space movements, we re-calibrate the robot to the grid to avoid position drift.
We collected 6 hours (14000 data samples) of the robot arm following uniformly random policy. There were episodes state resets. In addition to the robot, there several distracting elements in the image. A looped video htttp:://www.youtubbe youtube. wwtp//atcaaa==aaazzziiiaaa
plays on a large display in high resolution (4K video) at 2x speed. Four drinking loy birds, a color-changing lamp, and flashing streamer lights are also present. During the last half hour of image collection, the distracting elements are moved and/or or removed to simulate additional uncortainty in the environment. An illustration of the setup is Figure 9 along with the specific counter- exammpe for one-step inverse models. The quantitative performance of various basclines and AC-State this task is shown in Table 5
23

---

Under review as submission to TMLR
Autoencoder
One-Step Inverse Model
Generative Transition
AC-State (ours)
Figure 8: For four different baseline models we show input image (left), top-down image (center), and reconstruction of the image from the learned latent stale (right). Two consecutive frames are shown for each method. Both the one-step inverse model and AC-State successfully discard the background distractors, but only AC-State does this while also successfully capturing the true position of the robot arm.
Additional Result Videos of the latent state visualization for baselines and AC-State are shown in Figure [ In each video, the frontal view (ground truth) is shown on the left, the top-down view (ground truth) is shown in the middle, the reconstruction of the frontal view from the latent state is shown the right.
Algorithms
IMAGE REcONSTRUCTION ERBOR
GRouND TRUTH STATE ACCURACY
Dynamice Difference Error (%)
Contrustive
45.20
33.05
66.1
Autoeucoder
50.56
Forward Generative
65.48
21.30
76.1
1-Step Inverse
88.17 68.70
86.75
16.6
AC-State
88.04
99.83
0.52
Table 5: Robot arm results with various baselines. We show error in reconstructing the original image (generally higher is better it indicates that exogenous noise is discarded from the latent state). We also train probe network to predict the ground truth endogenous state from the learned latent state as well dynamics difference error from the learned latent dynamics.
24

---

Under review as submission to TMLR
1,0
2.0
0,1 M 2,
0.2 1.2 2.2
1,2
22
Figure 9: The robot arm has five actions and moves within nine possible control-endogenous states (left). The transition directions are indicated by the arrows. For example, if the robot arm is (0,0) and selects the down action, moves (0.1), but if selects the up action, remains at the same position. A simple inverse model can achieve perfect accuracy if the middle of true control-endogenous states are mapped to single latent state (middle). This leads the one-slep inverse model to merge them (right).
A.4 Matterport Simulator with Exogenous Observations
We evaluated AC-State the matterport simulator introduced in Chang al. (2017). The simulator contains indoor houses in which an agent navigate. The house contains finite number of viewpoints which the agent can navigate to. At each viewpoint, the agent has control of its viewing angle (by lurning left or right by an angle) and its elevation: in total there are 12 possible viewing angles per viewpoint and 3 possible clevations. We collect data using a random rollout policy. At cach stop of the rollout policy, the agent navigates to a neighbouring viewpoint. We also randomiy change the agent elevation al some of steps of the rollout policy, in order introduce exogenous information which the agent cannot control. We collect single long episode of 20,000 state-transitions. The condrol-endogenous latent state in this setup is the viewpoint information while the exogenous information is the information regarding agent elevation.
Experimental Details: The model input is the panorama of the current viewpoint ie.. 12 images for the 12 possible views of each viewpoint. The AC-State model is parameterized using a vision transformer (ViT) Dosovitskiy et al. (2021). Each view within the panorama is fed separately into the ViT as a sequence of patches along with a learnable token called the class (or CTSS) following the procedure in Dosovitskiy iet al. 2021p. To obtain the viewpoint representation, we take the representation corresponding to the CLS token of each view and take the mean across all views. We discretize this representation using VQ-VAE bottleneck van den Oord et al. (2017) to obtain the final representation. We use a 6-layer transformer with feedforward network (FFN) after every attention operation in the ViT similar to Vaswani et al. (2017). The FFN is 2 layer MLP with GELU activation Hendrycks & Gimpel (2016) which first projects the input to higher dimension D and then projects it back to the original dimension. We set the FFN dimension D to 512. We use 4 heads in the VIT. We train the model for 20 epochs using Adam optimizer Diederik et al. (2014) with learning rate 1c-4. The model is trained to predict the viewpoint of the next state as the action.
256 dimensions in the embedding. We
Baselines We use baselines for comparison Auto-Encoder In this baseline, we feed the observation as input ViT encoder and VIT decoder is trained to reconstruct the observation. Generative In this baseline, the model is trained to predict the next observation given the current observation and the action. The encoder and decoder both are implemented using a ViT model. Contrastive In this baseline, we use the simelr objective (Chen et al. 2020 We first pass the observation through a ViT encoder. We take the representation corresponding to the CLS token and split into two portions. We then apply the simelr objective considering the two representations from the same observation as the positive pairs and other combinations as the negative pairs. We also tried a temporal contrastive objective instead of simelr, but found that simclr had better performance in ignoring exogenous noise.
(1) Auto-Encoder, (2) Generative, and (3) Contrastive.
25

---

Under review as submission to TMLR
We present the results for this experiment in Figure (right). The Controllable Latent State Accuracy is the viewpoint, prediction accuracy for the current state. The Erogenous Noise-Ignoring Accuracy. is calculate as 1 8-33.33 where & is elevation prediction accuracy. Thus a higher elevation prediction accuracy leads to a lower the exogenous noisc-inducing accuracy. We can see that the proposed AC-State model has the highest. control-endogennous latent. state and exogenous noise-ignoring accuracy. Thus, outperforms the baselines we considered at capturing Control-endogenous Latent State information while ignoring exogenous noise. We calculated state parsimony Num. Ground Truth States Therefore, lower Num. Discovered States state parsimony denotes high number of discovered states which means that the model fails at ignoring exogenous information. The proposed model has the highest state parsimony which shows the effectiveness of the model in ignoring the exogenous noise whilst only capturing control-endogenous latent state. Detailed results with mean and standard deviation over three seeds are reported in Table 0
Results:
100-33.38
Algorithms
Endogenous Accuracy
Exogenous loiso-ign Accuracy (%)
Autoencoder
86.15 + 1.33
40.08 + 1.24
Contrastive
93.57 1.07
10.05 1
Forward Generative
97.25 1 0.38
3.66 + 1.19
AC-State (No Bottleneck)
96.94 + 0.58
47.81 = 2.77
AC-State
0.33
91.34 + 1.46
Table 6: Experimental Results with Matterport with mean and standard deviation over three seeds.
B Additional Discussion on Related Works
In this section, we provide a clear example of why an onc-step inverse model might not be useful to ignore exogenous noise. The idea of using a simple one step inverse dynamics models have been explored in the past (Pahhak et al.) 2017 Bharadhwaj et al.) 2022), yet the one step inverse model has counterexamples establishing that it fails to capture the full control-cndogane
latent state (Efroni et al.j 2022e Misra et al. 2020 Rakelly et al. 2021 Hutter & Hansen 2022p. Intuitively, the 1-step inverse model under- constrained and thus may incorrectly merge distinct states which are far apart in the MDP but have similar 0
local structure. As simple example, suppose we have cycle of states: 81, $2, 83, 84, 85, 86 where moves carlier in the cycle and = moves later in the cycle. Suppose $1, are merged into distinct latent state Si, 8k- The inverso-model examples are: (8,, 1), (8j, 8;, 0), (sj, Sk: 1), (sk, 0), (Sk, 8t, 1), (8t, 8k, 0). Because of these examples have distinct inputs, a 1-step inverse model still has zero error despite the incorrect merger of {s1, s1), {s2, ss), and 86).
are merged into a distinct latent state and 83, 86 are merged into a distinct latent state
C Algorithm Details
We provide more detailed algorithm descriptions for AC-State with a planning policy (Algorithm 2 The latter uses findgoal function which selects a low-count state using breadth-first search along with a plan function that finds optimal action using Dijkstra's algorithm.
26

---

Under review as submission to TMLR
Algorithm 2 AC-State for Latent State Discovery using Planning Policy
Initialize replay buffer D. Initialize encoder So. Assune number of samples to collect T, a set of actions A.
31 U (u(x)),
(AA, lg 1, and T
3: for t = 1, ... T do
4:
Ft+1 ~ p( 0z)
5:
8; = -21)
1t+1 - -Jo (2(411)
Update tabular-MDP M with (st, ae st41). if == by then Pick a new goal
10:
te =
11:
tg == (s
12:
ag ~ U(A)
13:
Add T
14:
else
15:
a, := plan( 9;
16:
K2 lg -
17:
j T and K ~ U(j
18:
L Lac-state (fo.j. 2, a, k) + Bottteneck (fo, + Bottleueck (fottteeeck(fo...... (5+5)
19:
Update to minimize C. by gradient descent.
D Detailed Theory and Discussion
D.1 High-Level Overview of Theory
We present an asymptotic analysis of AC-State showing it recovers the control-ondogenous latent state encoder f.. The mathematical model we consider the deterministic Ex-BMDP. There, the transition model of the latent state decomposes into a control-endogenons latent state, which evolves deterministically, along with noise term-the uncontrol-endogenous portion of the state. The noise term may be an arbitrary lemporally correlated stochastic process. If the reward does not depend this noise, any optimal policy may be expressed in terms of this control-endogenous latent state. In this sense, the recovered control-endogennuu latent state is sufficient for achieving optimal behavior.
Intuitively, the Ex-BMDP is similar to video game, in which "game engine" takes player actions and keeps track of an internal game state (the control-endogenous state component), while the visuals and sound are rendered as a function of this compact game state. modern video game's core state often orders of magnitude smaller than the overall game.
The algorithm we propose for recovering the optimal control-endogennus latent state involves (i) action prediction term; and (ii) mutual information minimization term. The action prediction term forces the learned representation f(z) capture information about the dynamics of the system. At the same time, this representation for fxx (which optimal for cction-prediction) may also capture information which is unnecessary for control. In our analysis we assume that f(x) has discrete values and show the control- endogenous latent state is the unique coarsest solution.
To enable more widespread adoption in deep learning applications, we can generalize this notion coarseness minimizing mutual information between æ and f(x). These are related by the data-procese prote-ppocssingg inequality; coarser representation reduces mutual information with the input. Similarly, the notion of mutual information is general it does not require discrete representation.
27

---

Under review as submission to TMLR
D.2 The Control-Endogenouu Partition is Bayes' Optimal Solution
Consider the generative process in which is sampled from distribution μ4, the agent executes policy a for time steps and samples x'. Denote by Pr,, (2,, () as the joint probability, and by (a æ, t) the probability that under this generative process the action upon observing æ is a. The following proof for Proposition 5.1. which builds on Proposition 3.3 shows that the optimal bayes solution Fx,,( (a x, 2',", t) is equal to (a f. for P. (x, (x)) O, where Px, x', is the probability to sample æ
Proposition 5.1. Assume that a is on endogenous policy. Let x ~ P for some distribution μa. Then, the Bayes' oplimal prediclor of the pcrion-prediction model is piece-wise constant with respect to the control- endogenous partition: for all a E A, t 0 and I, x' E X such that IP ( x' it > 0 it holds that
P (a x, 2,,',tt .t) = (a, (z))
Proof. Assume that is and endogenous policy. Let ~ p for some distribution μ.. Then, the Bayes' optimal predictor of the action-predictio model is piece-wise constant with respect the control-endogenous partition: for all A, t 0 and x, x' X such that (x, 0 it holds that:
P Pr.a(a x, 2,,2 Pr (a (.')),
We comment that the condition (=,,( (x, x²,t) > 0 is necessary since ,otherwise, the conditional probability (r, [o]] x4,x", well not defined.
Proposition 5.1 is readily proved via the factorization of the future observation distribution to control- endogenous and exogenous parts that holds when the executed policy does not depend on the exogenous state (Propasition 3.3
Proof. The proof follows by applying Bayes' theorem, Proposition equation 3.3 and eliminating terms from the numerator and denominator.
Fix any > x E X' and a e A such that (x', > Let s fz(x,, s f((x',, fx,((r), and e' (x)). The following relations hold.
Pa I1,
(o) Pr.a(si (I' Ir, a. ax,,,t)PP.. ( x))
Enr Paple x, a',
(6)
Prpp(ee x, ()
En Pripee (x' |x, a', (''() s)
(c)
g(x'1 s s',e')lPx |s, a, t) lele le, (e,t)((a (s) |s,a', (e' le, (,()a(dd 8)
En a106 18 (',e')D=,,ll
Pri s, a, () s)
En. Pana(ss a t)mm(a')
Relation (a) holds by Bayes' theorem. Relation (b) holds by the assumption that is endogenous. Relation (c) holds by Proposition equation 3.3
Thus, (a | x²,2, t) (a I f.)), and is constant upon changing the observation while fixing the control-endogemous latent state.
D.3 Stationary Assumption
We make the following assumptions the policy by which the data is collected.
28

---

Under review as submission to TMLR
Assumption D.1. Let TD(s' ss be the Markov chain induced on the control-endogenous state space by execuling the policy aD by which AC-Slate collects the data.
1. The Markov chain TD has stationary distribution pp such that up(s, a) > 0 and TD(a s) 2 Tmin for all & E S and a E A.A.
2. The policy TD by which the data is collected reaches all accessible states from any states. For any 8, s E and any > if 8' is reachable from 8 then PD(s' 8, h)) > O.
3. The policy TD does not depend on the erogenous state, and is an endogenous policy.
See Levin & Peres 2017), Chapter 1, for further discussion on the classes of Markov chains for which the assumption on the stationary distribution hold. The second and third assumptions are satisfied for the the random policy that simply executes random actions.
We consider the stochastic process in which an observation is sampled from distribution j such that p(s) = (s((s) for all Then, the agent executes the policy up for time steps. For brevity, we denote the probability measure induced by this process as PD.
D.4 The Coarsest Partition is the Control-Endogenouu State Partition
Proposition 5.1 from previous section shows that the multi-step action-prediction model is piece-wise constant with respect to the partition induced by the control-endogenous states f. f [S]. In this section, we assume that the executed policy is an endogenous policy that induces sufficient exploration. With this, we prove that there is no coarser partition the observation space such that the set of inverse models are piece-wise constant with respect to it.
We begin by defining several useful notions. We denote the of reachable control--ndogenous states from in h time steps as (s)
Definition 5.2. (Reachable Control-endogenoni States). Let the set of reachable control-endogennus states from S in h > 0 time steps be R(s, /t) {s' max, (s |s, s) [1)==]] 1}.
Observe that every reachable state from s in h time steps satisfies that max, (s. | so = h, due to the deterministic assumption of the control-endogenous dynamics.
Next, we define a notion consistent partition with respect to a set of function values. Intuitively, a partition of space X' is consistent with a set of function values the function pieco-wise wisc constant on that partition.
Definition 5.3. (Consistent Partition with respect to G). Consider set G {g(a, V') Eyee where A x y x y []. We say that f y + [N] is consistent partition with respect to G if for all V. 1/2 E , f(1)) f(12) implies that g(a, v. vi) g(a y, v/2) for all a E A.
Obscrve that Proposition 5.1 shows that the partition X' according to f. is consistent with respect to (Po(a] x, x', h)) x, E X. [H] s.t. Po(x, x',, > 0}, since, by Proposition 5.1 Pp(a x, h) = Pa(a
f+
Towards establishing that the coarsest abstraction according to the AC-State objective S. we make the following definition.
Definition 5.4. (The Generalized Inverse Dynamics Set AC(s,h)) Let S, E N. We denole by h) as the set of multi-step inverse models accessible from in h time steps. Formally,
AC(s, h) (Po(a s', s'"" h') R(s h'), E RRss,,''), A, [h]].
(5)
Observe that in equation equation the inverse function Po(a 8,88", is always well defined since Pp(ss' s", 16) > O. It holds that
ppls' Po(s')Po(s 1,11) 0,
29

---

Under review as submission to TMLR
since PD((') 0 and PD(s" | 1)) > 0. The incquality PD(s') 0 holds by the assumption that the stationary distribution when following U has positive support on all control-endogemous states (Assump- tion D.I.. The inequality PD(s" s', h/) > 0 holds since, by definition s" E R(s, h²) is reachable from s in ' time steps; hence, Ps(s" ss,hh') > 0 by the fact that Assumption D.1 implics that the policy Tp induces sulficient exploration.
Theorem 5.5. (f. is the coarsest partition consistent with respect AC-State objective). Given the bownded diameter Assumption 3.2 and the high-coverage policy assumption, (Appendir D.3) hold, it follows that there is coorser partition than fa, which is consistent with AC(s, D) for any E ES.
Assume 3.2 and D.1 holds. Then there is no coarser partition than f. which is consistent with AC(s, D) for any E S.
Proof. We will show inductively that for any h > 0 and 8 E S there no coarser partition than R(s, h) for the set R(s, h) that is consistent with AC(s, h). Since the set of reachable states = D time steps is S-all states are reachable from any state in D time steps-it it will directly imply that there is no coarser partition than R(s, D) S consistent AC(s, D).
Base case, h 1. Assume that h and fix some E S. Since the control-endogenous dynamics is deterministic, there are A reachable states from s. Observe the inverse dynamics for any s' E R(s, satisfies that
1 if leads sf from PD(a 8,8
(6)
0 o.w.
This be proved by an application of Bayes' rule:
FD(s' 8= = a, (a s)
PD(a 8,86, 1)
En Pp(s' | 8, s, 1)
T(s' s,a)To(a ()
En T(8 8, 8.
2 Tmin (s,a) leads to s
=0
O.W.
where the last relation holds by Assumption D.1 reachable from 8, the probability function PD(a a 8, 8,8, 1) well defined.
Furthermore, observe that since s' E R(s, 1), i.e., is
Hence, by equation equation we get that for any $1 s/2 E R(s,, 1) such that of + s/2 it holds that exists A such that
miin 2 Pa 8, s1 1) * (a $2 1) = 0.
Specifically, choose a such that taking a from leads to s and see that, by equation equation
Tmin 2 PD(a 8 1) # Pa 8, 2, 1) 0.
Lastly, by the fact that E S is an arbitrary state, the induction base case is proved for all E S. Induction step. Assume the induction claim holds for all f E [h] where h E N. We now prove it holds for t=h+1.
Fix some & E S. We prove the induction step and show that R(s, h + 1) the coarsest partition which is consistent AC(s, h + 1). Meaning, there exists 3, E [/ + 1],a a such that P(a st, P), 3, (,111) e AC(s, h + 1) and
Po(a 8, 8,61, + Po(a 8,62,
(7)
Observe that, by Definition 5.4 it holds that,
AC(s, 1) (Po(a s' h + 1)) UsER(N.) AC(s,h).
30

---

Under review as submission to TMLR
Meaning, the set AC(s, & +)) can be written as the union of (1) the set (Pv(a s, h 1() and (2) the union of the sets AC( h) for all 8 which is reachable from 8 in single time step.
By the induction hypothesis, the coarsest partition which is consistent with AC(3, h) is RR,, We only need to prove, that for any 81, 82 E R(ss h + 1) such that 81 # $2 exists some a E A, E [h] and 8he E RRss h + 1) such that
Pp(a Sw, 11,,),,1)) + PD(a sw, 1'),
this will imply that the set of reachable states in h + time states is also the coarsest partition which is consistent with As(s,
Fix $1, 82 E R(s, + 1) such that š1 # 82 we show that exists a certificate in AC 1) that diffcrentiate between the two by considering three cases.
1. Case 1: Both 81 and 82 are reachable from all 8 E R(s,1) 1). In this case, for all s' E R(s 1) it holds that 31, 82 E R(s',h) h). By the induction hypothesis, 81 and 82 cannot be merged while being consistent with AC(s',
2. Case 2: Exists s E R(ss 1) such & is reachable from s' in h time steps and &2 is not. Let be the action that leads to s' from state 8. In that case, it holds by the third assumption of Assumption D.1 that
Po(a s, 8,81, 1) (a)
PD(s1 s, h. + 1)Tp(a (s)
Po(s1 s, 1)) ()
PD (), s)
Ea Pp(51 s, h 1)
PD(51 s',h)
Mmin En. PD
8, h
0.
(8)
Relation (a) holds by Bayes' rule. Relation (b) holds by the fact that (s.a) determinstically leads to s'. Relation (c) and (d) holds by Assumption D.1
Observe that Pa 8, 5,, h 1) 1) 0 since 52 is not reachable upon taking action from state by the assumption. Combining this fact with equation equation implics that
0 < P>(a 8, 8, h + 1) + Po 8, 52, h 1) =.
Hence, exists certificate that differentiates between 5n and 82- Observe that since E R(s, h h 1); ie., it is reachable from 8, it holds that Pa(a s, f + 1) 0, i.e., it is well defined.
3. Case 3: Exists s' E R(s, 1) such 82 is reachable from s' in h time steps and 81 is not. Symmetric to case
This establishes the result we needed to show in equation equation and, hence, the induction and result hold.
E Discussion: What does Control-Endogennuu Latent State Capture? Does it ignore Dask-Relevant and Reward-Relevand Information?
The control-endogenous latent state defined by having factorized dynamics that depend on actions, while the exogenous dynamics do not depend on actions. If we refer to the control-endogenous state as and the exogenous state as e, we can write the factorized dynamics as: P(s'|s, a)P(cl In particular, we want find the smallest. such that the latent dynamics follow this factorization.
31

---

Under review as submission to TMLR
Causal Dynamics Learning (Wang et al. 2022b provided some definitions that are helpful for building intuitions about the semantics of the control-endogenous state. In their work, they decompose factorized state into controllable factors, action-relevant factors, and action-irrelevant factors. They define controllable factors as any causal descendants of actions or other controllable factors. They define action-relevant factors as causal parents of controllable factors or other act ion-reeelln factors. All other factors are defined as action-irrelevatt. While AC-State learns a. general encoder instead of relying on a given factorization, the special case of known factors can be used to build intuition.
An example is given in Wang et al. 2022b involving six known factors: 21, 22, 23, 21, 20, 26. The dynamics are factorized as: P(z a))((= 21, 22, 23 ((sslss, =4 (=4))(2()))((( Rs. 26 (=))P((((22)) (s)) We can assign these factors to either the control-endogenous state or exogenous state. The smallest way to define the control-endogenous state is: s = (z1, 22, 23, 21), while the exogenous state is set to e (2). Intriguingly, the control- endogenous state includes many factors which are not controllable.
As simple example, we can imagine a robotic arm that is allowed to interact with two blocks (one red and one blue). What is the eondrol-endogenous latent state for this environment? It clearly includes the robotic arm itself, as it can be manipulated using the actions. Likewise, the physical properties of the blocks (their position and weight) will be included in the control-andogenons latent state.
The color of the blocks is not part of the control-endogenous latent state. What if our task is to retrieve the red block? In general, the control-ndoogeoous latent state will not be sufficient for solving the task. If task completion has causal effect on the rest of the control-endogenous state, then the task relevant information is control-endogenouu. As a concrete example, if the agent were forced to redo the task of retrieving the rod block until successful, then the color of the blocks would become control-cndogenous. Intuitively, in the most natural and realistic environments, rewards and task completion will have causal effect on the rest of the control-endogenone state. On the other hand, if the episodes are terminated immediately after reward is received, it would be useful to add reward prediction as an auxiliary task for learning the representation.
E.1 AC-State Captures the Position of a Block that is Movable by the Agent
We also ran experiments to verify that AC-State performs correctly in an environment with an agent and a movable block. According to theory, both the agent and the movable block should be captured by the representation.
we write the agent's position as SA and the block's position as 8B, the ground truth dynamics are P(('A sgls,, 8B. = ls (s a). Note the asymmetry: the agent affects the block but the block does not affect the agent. We tested two variants of this setup where the agent either pulls or pushes the block. In the pull environment, the block is pulled alongside the agent unless it moves to the outer edge of the environment, in which case the agent drops the block. In the push environment, the agent moves the block if it moves into it, and the agent/block /look wraps around to the other side of the grid (like a lorus) if it moves off the edge. These setups are designed so that the block cannot be stuck in a corner of the grid.
AC-State was trained on 500k samples in each environment with a uniformly random rollout policy. Prob- ing revealed that both S and sg were perfectly captured by the encoded representation. Planning was performed using Dijkstra's algorithm and example rollouts are included in the supplementary material pull examples. txt and push examples.txt.
32


Reviews:
[{"review_id": "DeuimWEIEF", "reviewer": "Reviewer_SvYL", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- this work provides both intuitive explanations and a theoretical analysis of the method;\n- the experiments provided in this work is neatly designed to showcase the problem setting where AC-state has an advantage over other methods; the maze environment is very interesting;\n\nWeakness:\n- lacking discussion of some related work in causal learning; the idea of learning agent-controllable states is very much related to the literature of learning causal models for robots (be it policies, dynamics models or state representation); prior work has explored this front a lot (to list a few representative ones see ref. [1-3]) and it is important to distinguish how AC-state is different from ideas in causal learning\n- The empirical results are limited to learning reachable states and no downstream control task performance is shown; is the method limited to only these tasks because the assumptions of the AC-state are too strict? How would someone leverage AC-state for better downstream task learning?\n- Is AC-state effectively learning to ignore noise in the environment? What if the noise is biased/spurious given the control? would AC-state still be able to recover the true invariant state? \n- The implementation details are not sufficient for reproducing the results; what is your network architecture? What are the hyperparameters and how were they tuned? Was there any domain knowledge needed in tuning the model or do generic parameters work for a range of tasks?\n\n\n[1] Mastakouri, A. A., Schölkopf, B., & Janzing, D. (2021, July). Necessary and sufficient conditions for causal feature selection in time series with latent common causes. In International Conference on Machine Learning (pp. 7502-7511). PMLR.\n\n[2] Li, Y., Torralba, A., Anandkumar, A., Fox, D., & Garg, A. (2020). Causal discovery in physical systems from videos. Advances in Neural Information Processing Systems, 33, 9180-9192.\n\n[3] Fu, X., Yang, G., Agrawal, P., & Jaakkola, T. (2021, July). Learning task informed abstractions. In International Conference on Machine Learning (pp. 3480-3491). PMLR.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": " An agent interacting with an environment with dynamic background changes (noises) would need to separately model the controllable aspect of the state in order to correctly plan or learn robust policies. This work presents a method to learn agent-controllable latent states through learning a multi-step inverse dynamics model (i.e. predict actions given future observations).\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Related work section: Compare AC-state discovery with the literature of causal learning / counterfactual learning. How would methods in learning causal dynamics model perform in Ex-BMDP? Is AC-state discovery a simpler problem than causal inference? If so, in what scenarios would AC-state be preferred over causal methods?\n\n- Network architecture & Ablations: provide details of model design and loss terms used; show model performances ablating hyperparameters such as number of inverse dynamics steps and bottleneck state size;\n\n- Limitations and future work: discuss the limitations of AC-state or plans for scaling AC-state to complex robot learning problems (i.e., demonstrate superior performance of learning with the latent space discovered by AC-state on downstream imitation or reinforcement learning tasks).\n", "broader_impact_concerns": "There is no ethical concerns to be addressed.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "KoFGiV98Fa", "reviewer": "Reviewer_vkF8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n- Learning an encoder that can obatin control-endogenous latent states from rich sensory input is important to achieve efficient RL in the real-world, where the true state of an agent is often unknown. The paper proposes a simple yet effective algorithm based on a multi-step inverse model objective with information bottleneck, which is easy to train with supervised learning and makes the approach scalable.\n- The paper performs theoretical analysis to justify the multi-step inverse model objective.\n- Empircally, experiments on three different tasks show that the proposed approach is effective and outperforms the baselines.\n- The paper is well-written and the organization makes it easy to follow. The related work is thorough and clearly highlights the difference with prior work.\n\n**Weaknesses:**\n\n- The three tasks in the experiments use different baselines. To fully demonstrate the effectiveness of the approachs, the paper should compare with all baselines (especially stronger ones) used in the first-person navigation task on the other two tasks as well.\n- It is unclear how the Information bottleneck used in the encoder affects the theoretical analysis in Sec. 5.2.\n- The analysis assume determinstic dynamics, while in real-world scenarios, stochastic dynamics are common. The three tasks in the experiments also mainly have determinstic dynamics.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a simple yet effective approach to extract control-endogenous latent states (e.g., position of an agent) from complex observations and discard control-exogenous information (e.g., texture, lighting). The paper proposes a simple multi-step inverse model objective that predicts the first action given the current observation and a future observation. Information bottleneck is applied to the latent representation to achieve a minimal latent state. Theoretical analysis is performed to show that the minimal representation minimizing the multi-step inverse model objective is the control-endogenous state. Experiments are performed on robotic arms, maze exploration, and first-person indoor navigation, which demonstrate the proposed method’s superior performance over the baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- The claim “contains all of the information necessary for controlling the agent, while fully discarding all irrelevant information” is too strong. How do we know that the learned latent states “fully” discard all irrelevant information?\n- Consistent baselines across all experiments.\n- Discussion on how information bottleneck affects the theoretical analysis.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "UxBgmRxx85", "reviewer": "Reviewer_jojo", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n- The paper tackles an important problem, which is how to learn a low-dimensional representation from high-dimensional inputs while being invariant to exogenous noise.\n- The method comes with theoretical guarantees, which is nice.\n- The empirical evaluation is done across diverse environments: robotic control, multi-agent, and visually-rich navigation.\n\n### Weaknesses\n\n- I found the organization of the paper difficult to follow.\n- The theoretical analysis is dense and appears to be included at the expense of more intuitive explanations for why/how the method works.\n- There are many details missing about the experiments and how the method is implemented in practice.\n- The difference between this paper and Efroni et al. (2022c) is unclear.\n- The experiments do not provide insight into how the method compares to existing approaches that are also designed to ignore exogenous noise.\n- The assumptions needed for the method to work are quite strong and it is unclear how the method behaves if these assumptions are not met.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a representation learning algorithm, AC-State, which is designed to exactly recover the a discrete control-relevant latent space (the \"control-endogenous state\") while discarding other control-irrelevant information (the \"exogenous state\"). The algorithm relies on several assumptions: a discrete latent space, a bounded diameter (e.g. no absorbing states), deterministic dynamics, and that the data generated to learn the representation is generated by a endogenous policy (i.e. which is not sensitive to the exogenous state). The method works by learning an encoder $f(x)$ which maps observations to latents and an inverse k-step model in the latent space which predicts $a_t$ given $f(x_t)$ and $f(x_{t+k})$. It also incorporates an information bottleneck to restrict the capacity of $f$. The paper presents both a theoretical analysis of AC-State as well as an empirical one, demonstrating theoretically that AC-State will converge to the true control-endogenous latent space in the limit and empirically that it does so in practice. Empirical experiments are presented in three domains: a robotic arm with visual distractors (e.g. video playing in the background), a maze with distractor agents, and navigation in the Matterport environment. ", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The paper needs to be significantly reorganized and parts rewritten in order to present the ideas in a clearer fashion before I would be able to recommend acceptance. I have a number of specific suggestions for how this can be done, some of which I think are necessary and some of which I'd just strongly recommend.\n\nNecessary changes:\n- I did not find Figure 1 to be particularly helpful in understanding the method. It's not obvious which parts are learned or how they are learned (e.g. what the losses are). There is no mention of the information bottleneck which appears to be an important part of AC-State. The caption is also insufficient and should give a summary of the method.\n- Section 3 discusses a number of related approaches which attempt to discard exogenous noise, but dismisses them because they are not theoretically guaranteed to obtain the control-endogenous state. This is certainly a weakness of these approaches, but more explanation is needed regarding how AC-State differs from these approaches at a high-level. (Ideally, at least some of these methods would be included as baselines to demonstrate the situations in which AC-State's theoretical guarantees make a difference, though this is not strictly necessary for acceptance).\n- Table 1 lists several methods with no explanation, description, or citations (e.g. PPE, DBC, etc.). If they are important enough to be included in a table then they need to be named and discussed in the main text too. The table caption should also include citations for these methods.\n- More explanation is needed to clarify the difference between Efroni et al. (2022c). I understand that the current paper removes an assumption about the initial control-endogenous state, but in what other ways are the papers the same/different? Is the method otherwise pretty much the same? Does it have the same limitations regarding discrete state spaces, bounded diameter, etc.? Can it be included as a baseline in the experiments later in the paper?\n- Page 7 introduces the information bottleneck and states that it is important that the method finds the lowest-capacity f. This claim is not explained or justified. Please add an explanation and intuition for why this is needed.\n- The loss term $\\mathcal{L}_\\mathrm{Bottleneck}$ is not defined.\n- The overall algorithm including both loss terms and the architecture should be summarized somewhere (ideally in a figure).\n- The subsections on Page 8 (\"AC-State with Random Policy\" and \"AC-State with Planning Policy\") seem like they belong in the experiments section, as they are not really parts of the core method but rather approaches for data collection. Though that said, I was confused about how \"AC-State with Planning Policy\" works. Is the idea that you are given a planning policy, which is just used for data gathering to train AC-State? Or that it's used on top of AC-State (in which case, I would wonder how is the latent space learned)? Or is it used in-the-loop with AC-State (in which case, are there any stability issues with training, that often occur in deep RL)?\n- Section 5.1 seems to be from a previous paper. Is it necessary to reproduce here? It also feels a little bit like math for math's sake---isn't it kind of by definition that $P(a|x,x')=P(a|f(x),f(x')$ if $f(x)$ recovers the endogenous state and $a$ is determined solely by the endogenous state? Also, it seems that an important conclusion from this section is that the action distribution is piecewise constant, but the paper does not explain why this is the case (if I understand correctly, I guess it is because of the block-MDP structure?)\n- Several propositions from the appendix, e.g. Proposition 9, Assumption 10, etc., are referenced without any indication of what they say or entail. (There is also a typo on Page 9 which says \"Proposition 9 from the previous section\"). Based on the numbering it is also not obvious they are located in the appendix. Any proposition or assumption from the appendix that is referenced in the main text should be referenced along with a statement saying it is in the appendix and a summary of what it says. (One suggestion would be to use numbering like Proposition A.1, where the \"A\" indicates appendix).\n- I don't understand the point of Definitions 6 and 7 are, or what their relevance is for Theorem 8.\n- The experiments section needs considerably more detail, enough to know at least the high-level experimental setups. I appreciate some of the details are in the appendix but there is both not enough even in the appendix and definitely not enough in the main text. Here are some specific things that were unclear to me (though this section should include more details in general, not just the specifics I point out here):\n  - In Section 6.1, details need to be included regarding: the nine positions the arm can be in (i.e. describing Figure 3 in more detail), what policy is used to collect the 14,000 samples for training AC-State, and what the \"counterexamples\" mentioned in Figure 5 are.\n  - In Section 6.2, what policy is used to generate the data for the top half of Figure 4? And as discussed above, the planning policy needs to be more clearly explained.\n  - In Section 6.3, how are the various metrics (controllable state accuracy, exogenous noisy-ignoring accuracy) computed? How are the various baselines in Figure 6 implemented? e.g. what is the \"contrastive\" baseline---there are many contrastive methods, which one was used here?\n\nRecommended changes:\n- More intuition should be provided in the introduction about how/why AC-State works so that its relation to other methods can be better understood in the background section. I realize the introduction states *what* AC-State does---e.g. that it learns a multi-step inverse model with an information bottleneck---but it doesn't explain the motivation for this or *why* it works. The explanation right at the end of Page 7, for example, would be great to include earlier in the paper.\n- I would recommend swapping the order of sections 2 and 3, reducing the length of sections 2 and 5, and lengthening section 6. The figures should also be reordered to come in order of mention in the paper: figure 1, 3, 2, 5, 4, 6.\n- Definition 1, 5, 6, 7; Proposition 4; and Theorem 8 are never referenced again after being introduced. Definitions etc. should only be included if they are actually referenced. Please remove these or ensure they are properly referenced.\n- The paper would be much stronger if experiments were included demonstrating how well AC-State performs when its assumptions are violated. Does it degrade gracefully or catastrophically? What are its failure modes?", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers agreed that the motivation for the AC-state innovation presented in this work was clear and the combination of theoretical analysis and empirical investigation is compelling.  However, the reviewers were split as to whether the work is sufficiently clear to accept in its current form.  Reviewer jojo found the motivation and core contributions of the work compelling (the problem, the theory, and the experiments).  However, reviewer jojo also found the work poorly organized and a bit unclear throughout.  The reviewer identified a number of specific clarity issues.  The authors made an attempt to address some of the issues raised by the reviewer, which the reviewer acknowledged, but the final response of the reviewer is that the changes were not comprehensive enough.  The reviewer included some precise final points and emphasized their criterion for clarity being: "Is there enough detail in the main text of the paper that someone could implement the overall structure of the method and experiments...and is there enough detail in the appendix that they could get all the details too".  On the other hand, reviewer vkF8 found the paper "well-written and the organization makes it easy to follow".  The reviewer requested that the authors include consistent baseline algorithm comparisons for each of the 3 experimental settings, which the authors did.  However, consistent with the spirit of reviewer jojo's concerns, the presentation of the new experimental results does not have a consistent format, with very little results in the main text.  Reviewer SvYL also noted that implementation details in the initial submission are not sufficient for reproducing the results, and it was not clear if the authors addressed this specific point.

Summarizing the reviewer feedback, informed by reviewer jojo's points, I think the current draft needs something that could be either framed as minor revisions or a bit more than that, depending on the perspective.  Since the residual issues center on clarity and presentation, I do not think it is a good use of reviewer time for this paper to be re-reviewed separately to address this, so I propose to accept this with requested revisions, that I will check for before approving the final version.  I would encourage the authors to make as thorough a pass at revisions as they can, informed by reviewer jojo's additional feedback.  I'll produce a digest of a few specific points below that should be treated as a minimal set of required changes. 

Some specific areas to focus revisions around:
- The point being made by figure 2 could be clearer. Edit the figure and/or caption.
- Section 4 introduces AC-state quite generically, but parameter choices made for each experiment type are lacking (as noted by reviewer SvYL).  As a concrete example, for each of the three experiments, the horizon parameter should be discussed.
- The empirical results figures (currently figures 4, 6, and 7), are awkwardly placed in the paper (as noted by reviewer jojo), with figure 4 coming way too early and the order of maze and robot figures switched(?).  These figures seem like they should occur within section 6.  
- More clarity around the algorithm in the main text would be useful (include a version of the algorithm box, or equivalent information, in the main text).  
- Each of the three experiment types should have a consistent summary table with core metrics in the main text. As previously requested by the reviewers, the baselines are the same, but there is inconsistency in how and where the results are presented.
- Overall, aim to make the main text reasonably self-contained, with additional detail in the supplement.  As reviewer jojo pointed out, some rather essential information is presently in the supplement.
- Address jojo's remaining major points: algorithmic clarity around when updates happen relative to agent sampling, enumerate the baseline choices clearly in the main text, experimental details including number of seeds in the main text, etc.
Decision: Accept (poster)