Learning Action-based Representations Using
Invariance

Learning control for complex decision-making from high-dimensional observation spaces such as video and depth is vital for real-world applications of reinforcement learning (RL). To do this, a representation of the observation space allows agents to reason about the environment and take intelligent actions. However, learning these representations is often sample inefficient. One reason for this is that real-world scenarios often contain many irrelevant and distracting features embedded in a high-dimensional space. Correlating reward with relevant state elements, and not causally confusing distractors in this setting, is challenging—especially since reward signals are often sparse.

Representation learning has emerged as a promising approach to address this challenge by extracting a compressed and informative representation of the observation space that is useful for learning (Bengio et al., 2013). Representation learning removes irrelevant distractors from the state space used to learn the policy, which improves sample efficiency and performance. In RL, task-specific representation learning uses reward or expert behavioral similarity (Ferns et al., 2011; Agarwal et al., 2021) to discover the compressed representation, only describing task-specific elements. This has the advantage of capturing only information that is either useful for solving the task or relevant to the demonstrations while being limited by requiring either expert behavior or task-achieving policies, both of which can be difficult to obtain prior to learning. On the other hand, task-free methods use unsupervised signals like reconstruction (Lange & Riedmiller, 2010) and contrastive objectives (Laskin et al., 2020) and can be pre-trained on any data, including random actions. However, these methods are trained without action information. As a result they can capture exogenous distractors that are not useful for improving RL policy performance.

One promising direction of task-agnostic methods utilizes controllability to learn a behavior-relevant representation that is not task-specific (Lamb et al., 2022). These representations can avoid capturing task-irrelevant information while not requiring expert or reward-achieving behavior. Recent work in action-based representation learning for RL has shown promising results (Zhang et al., 2022) by utilizing inverse dynamics models to extract representations (Islam et al., 2022). These representations rely on a window of information by predicting the first action between two states separated by $k$-steps. If $k$ is small this representation is myopic, but when $k$ is large the prediction problem is underspecified. This underspecification restricts large $k$ to offline datasets with correlated action data—such as expert trajectories.

We investigate utilizing a novel invariant metric to learn a multi-step control-based representation instead of directly applying $k$-step prediction. Our action-bisimulation metric offers a novel framework for controllability metrics that takes a myopic dynamics encoding and extends it to multi-step representations. This formulation is inspired by reward bisimulation (Zhang et al., 2020b), which utilizes single-step reward information to learn multi-step return-capturing representations. Action-bisimulation applies bootstrapping on the myopic $k=1$ controllability representations to enforce multi-step invariance in an action-bisimulation encoding. Since the base case uses single-step prediction, the encoding can be trained with any offline data, even fully random. At the same time, boostrapping extends the action-bisimulation encoding to capture long-term controllability. LABEL:fig:bisim_main captures how action-bisimulation maps control-irrelevant states together, while not doing the same for control-relevant states.

This work offers an empirical analysis and theoretical formulation of the novel control-based invariant metric for representation learning. We demonstrate empirically that in scenarios where complex, long-horizon, sparse-reward decision-making is required, the metric improves sample efficiency compared to RL agents trained directly from pixels, or pre-trained with existing representation learning methods in multiple domains. Next, we provide qualitative results demonstrating the robustness of the learned representation to uncontrollable distractors, as well as sensitivity to control-relevant state features.

Representation Learning in RL. Learned representations have been widely applied to RL, formalized (Li et al., 2006) through hierarchical symbolic representations (Konidaris et al., 2014; Andre & Russell, 2002), skill abstractions (Dietterich, 2000), policy optimality (Auer et al., 2008; Jong & Stone, 2005; Abel et al., 2016), selective attention (Jones & Canas, 2010) and contingency awareness (Bellemare et al., 2012). One effective strategy is to use the representation to learn a model that is effective for planning (Hafner et al., 2019; Koul et al., 2023). These methods learn world models (Ha & Schmidhuber, 2018) and other representations that can be used for prediction (Singh et al., 2012), data generation and planning. Alternatively, other methods apply representation learning for filtering (Krishnan et al., 2015; Karl et al., 2016) or reduced complexity  (Higgins et al., 2016; Oord et al., 2018; Laskin et al., 2020) representations. Action-bisimulation is a novel encoder that learns controllability-based representations to improve RL performance. Unlike other representation learning methods, action-bisimulation uses a soft invariance pseudometric to capture action information through time.

Action-based Representations. RL methods have directly leveraged action-relevant representations in several ways. This includes contingency awareness (Bellemare et al., 2012; Choi et al., 2018; Chuck et al., 2020; 2023), which is closely related to action controllability (Zhong et al., 2020) and control information measures like empowerment (channel capacity between actions and state) (Jung et al., 2011; Mohamed & Jimenez Rezende, 2015; Levy et al., 2023) or affordances (Cruz et al., 2016; Khetarpal et al., 2020; Nagarajan et al., 2020). Multi-step inverse models are most similar to action-bisimulation, but common multi-step inverse methods (Lamb et al., 2022; Islam et al., 2022; Koul et al., 2023) require selecting a specific $k$ for the multi-step horizon, potentially leaving critical control information on the table. Further, it has been shown that multi-step inverse models can be insufficient when the dynamics are periodic (Levine et al., 2024). Action-bisimulation uses a soft invariance metric to extend single-step models, which better preserves long-term controllability.

Bisimulation methods. Bisimulation describes future invariant state representations, originally applied to stationary representations (Larsen & Skou, 1989; Dean et al., 1997; Ferns et al., 2004), before being extended to continuous state MDPs (Ferns et al., 2011). Reward-based bisimulation methods have gained popularity through learned deep representations (Zhang et al., 2020b). This has been extended to non-optimal policies (Castro et al., 2021), with generalized value function bounds (Kemertas & Aumentado-Armstrong, 2021) and augmented with state discretization (Kemertas & Jepson, 2022) and clustering (Liu et al., 2023). Bisimulation-based methods have also been applied in different contexts: expert policy similarity (Agarwal et al., 2021; Bertran et al., 2022; Mazoure et al., 2021), goal-conditioned RL (Hansen-Estruch et al., 2022) and reward-action policy equivalence (Liao et al., 2023; Castro, 2020). While this work draws on reward-bisimulation, action-bisimulation is fundamentally offline and task-agnostic because it takes an expectation over actions, removing its dependence on any policy.

A Markov decision process is defined by the tuple $\mathcal{M}\coloneqq(\mathcal{S},\mathcal{A},p,R)$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space and $s\in\mathcal{S},a\in\mathcal{A}$ are states and actions respectively. $p(s^{\prime}|s,a)$ is the transition function that gives the probability of the next state $s^{\prime}$ given the current state and action $(s,a)$. The reward function $R(s,a)$ maps state and action to a scalar reward. A policy $\pi(a|s)$ is the probability of an action given the current state.

This work utilizes the following two-phase paradigm: in the first phase, the agent first takes actions without access to the reward function $R(s,a)$ to generate a dataset of ordered state action tuples $\mathcal{D}\coloneqq\{(s^{(0)},a^{(0)}),\ldots(s^{(|\mathcal{D}|-1)},a^{(|\mathcal{D}|-1)})\}$. Then, a representation $\phi:\mathcal{S}\rightarrow\mathcal{Z}$ is learned from $\mathcal{S}$. In the second phase, the agent learns from extrinsic reward utilizing the learned representation.

Bisimulation Metrics (Ferns et al., 2011; Castro, 2020) soften the notion of state partitions with a pseudometric space $(\mathcal{S},d)$, where distance function $d:\mathcal{S}\times\mathcal{S}\rightarrow R_{\geq 0}$ measures the similarity between two states.¹¹1This is a pseudometric, meaning that two different states can have $0$ distance. The on-policy bisimulation metric (Kemertas & Aumentado-Armstrong, 2021) is:

where $W_{1}(d)$ is the 1-Wasserstein distance and $c$ is a scalar hyperparameter that weights the multi-step sensitivity of the distance. The 1-Wasserstein metric measures the distance between next-state distributions in the latent bisimulation space. We propose a novel controllability-based relation, which replaces reward equivalence with single-step control equivalence. By replacing rewards in the equivalence, the relation is task-agnostic.

This equivalence can be similarly relaxed into a pseudometric. However, in the off-policy setting, we are not interested in a particular policy, but all policies. Thus, action-bisimulation uses the expectation over uniform actions to encode all possible policies.

In the next section, we describe how $\psi(\cdot)$ can be learned from data, and how to use $d_{\text{a-bisim}}$ to learn an action-bisimulation encoder.

This section describes the algorithm for training an action-bisimulation encoder. First, the single-step encoder is learned, then the distance in single step space is used as the “base case” for the recursive step. The training flow and inputs are visualized in LABEL:fig:bisim_mainb.

In this section, we aim to answer the following questions: 1) Does pre-training with the action-bisimulation objective learn representations useful for arbitrary downstream tasks? 2) How does this pretraining compare with existing methods, especially single-step action controllability? (3) Are the learned representations robust to background distractors? (4) How well does the action-bisimulation procedure capture multi-step relationships between state elements?

We evaluate experiments in three domains illustrated in Figure LABEL:fig:downstream_rl. Nav2D is a 15x15 grid environment where the agent navigates using cardinal directions to the center of the grid, avoiding randomly generated 2x2 obstacles. Pointmaze (Fu et al., 2020) is a 2D Mujoco control environment where the agent takes actions to reach a goal location while being impeded by obstacles. We also investigate Distractor Pointmaze where the background in Pointmaze has been replaced with photorealistic distractions in the form of video clips. Finally, Habitat (Savva et al., 2019b) is a complex 3D environment where the agent must navigate through scans of human environments to reach a goal location. Additional environment details are in Appendix H (number of obstacles, goal/grid size, randomization, etc.) and all other relevant hyperparameters are in Appendix I.

Controllability-capturing encodings for reinforcement learning are a promising direction for representation pretraining since they can be learned without reward but are still able to filter out uncontrollable distractors. However, existing methods either only capture short-term controllability or are dependent on demonstration data, which has implicit task bias. We introduce the action-bisimulation encoding, which builds off of myopic representations by enforcing recursive invariance to learn a supervision-free multi-step controllability representation. The empirical results in this work demonstrate how these encodings can be used to improve the sample efficiency, especially in domains with significant background distractors. The primary limitation of this method is the inverse dynamics single-step model, which might not capture all controllable features, just a subset. This can result in the representation being agnostic to important task elements. A more in-depth discussion of limitations is included in Appendix F. Altogether, action-bisimulation is a novel invariance relation for capturing controllability from offline data that removes expert performance requirements and smoothly handles long-horizon controllability.

This work has taken place in part in the Safe, Correct, and Aligned Learning and Robotics Lab (SCALAR) at The University of Massachusetts Amherst. SCALAR research is supported in part by the NSF (IIS-2323384), AFOSR (FA9550-20-1-0077), and the Center for AI Safety (CAIS). The work was supported by the National Defense Science & Engineering Graduate (NDSEG) Fellowship sponsored by the Air Force Office of Science and Research (AFOSR). Special thanks to collaborators Stephen Guigere, Harshit Sikchi, Alex Levine, Siddhant Agarwal, Rudolf Lioutikov, Yuchen Cui, Akanksha Saran, Wonjoon Goo, Daniel Brown, Prasoon Goyal, Christina Yuan, and Ajinkya Jain for their fruitful conversations and timely help.