Measuring Empowerment from Offline Data for Intrinsically Motivated Exploration

Introduction

Online exploration is a challenging yet central problem in reinforcement learning, where the agent must traverse the environment to seek out novel states with potentially high reward values. Previous methods tend to place equal weights across all states in the prior, where the agent is optimized to visit all of these states until all of them have been visited. But without any prior, the agent must visit all states, increasing the sample-complexity. This is suboptimal as some states naturally grant the agent more control over what it could do (e.g., picking up a key lets the agent open a door to a new room).

A natural way to characterize the idea above is to consider empowerment [19], a metric that measures to what degree an agent can influence its future states via its actions. States with high empowerment have many distinct future outcomes that the agent can achieve by taking distinct sequences of actions. Because these high-empowerment states have more future outcomes, it seems reasonable to explore more around these states as an agent has the potential to discover more future outcomes or behaviors. Conversely, in low-empowerment states, there is very little that an agent can do to influence its future outcome, so it should spend less time exploring these regions of the state space as there is less strategic value to be gained.

Previous works have explored using empowerment, or related information-theoretic objectives for online exploration [33, 18, 22, 5, 7], but operate in the purely online setting, where we do not have access to prior data. Luckily, in many practical environments, we have prior data in the form of unlabeled trajectories. Ideally, we would like to use the prior data to inform us how to explore online, but estimating empowerment purely from offline data is difficult: existing empowerment estimators (Mohamed & Jimenez Rezende, 2015; Gregor et al., 2017; Eysenbach et al., 2019; Sharma et al., 2020b;a; Kumar, 2018; Becker-Ehmck et al., 2021; Bharadhwaj et al., 2022; Kim et al., 2019) rely on online samples where a skill-conditioned agent must interact with the environment during the optimization. Furthermore, adapting online methods to the offline setting by simply regularizing the policy as is typical in the offline RL literature [21, 13, 20] does not work (Section 6.2). Existing empowerment estimators look at whether the outcomes achieved by a policy align with its intentions; however, offline, we do not know the policy’s intention. How can we learn empowerment from offline data to improve online exploration?

In this work, we provide a novel algorithm for learning empowerment purely from offline data, without additional environment interaction. Our key contribution is an offline estimator of empowerment: by learning an empowerment-maximizing policy in tandem with its discounted state occupancy, we can extract an estimator of the empowerment from prior data alone. We show that our method produces qualitatively accurate empowerment maps, correctly identifying which states admit the most diverse future outcomes. We further demonstrate that the offline-extracted empowerment estimate can be used to improve exploration in online tasks. Specifically, we find that using the empowerment as an auxiliary loss term improves exploration on downstream, online tasks in OGBench environments [29]. We evaluate our methods in sparse-reward, goal-conditioned environments, where knowing where to explore is harder due to the sparsity of the reward, and the agent must prioritize which states to visit. In this setting, our empowerment estimator provides a principled way to characterize the exploration value of a state from offline data, directing the agent toward regions with greater potential for diverse future behaviors.

Related Work

Empowerment. Empowerment [19] measures an agent’s control over its future as the mutual information between its action sequence (or skill choice) and the resulting future state. Prior works have estimated empowerment from online data by optimizing a variational lower bound [26], and later work estimates the empowerment of an agent’s current policy from off-policy samples, by maximizing lower bounds on the empowerment [22, 6, 27, 7, 9] (Zhao et al. 2021; Zheng et al. 2025). However, these methods estimate the empowerment of the current goal-reaching policy rather than the maximum achievable empowerment, which can be a very weak bound in practice [31]. Becker-Ehmck et al. [5] compute empowerment directly, but require finite action spaces and incur cost exponential in the horizon. Schneider et al. [32] show that empowerment pretraining provides a strong initialization for downstream finetuning, but use a fixed estimator constructed by a human annotator. The closest to our work is EELMA (Song et al., 2026) which estimates the effective empowerment (defined in § 3) using contrastive learning for discrete LLM actions from collected trajectories; we adapt this method to the continuous action space setting and use it as a baseline in § 6.2. Critically, none of these empowerment estimators work well when adapted to the offline RL setting (Section 6.2). Our work is the first to estimate the maximal empowerment purely from unlabeled offline data.

Intrinsic motivation and curiosity-driven exploration. A large body of work in reinforcement learning has studied intrinsic motivation as a mechanism for encouraging exploration in sparse-reward environments. Early approaches include count-based exploration and density modeling methods that assign bonuses to rarely visited states (Tang et al., 2017). More recent methods such as curiosity-driven learning estimate prediction error or learning progress in a forward dynamics model as an intrinsic reward signal [8, 18]. Our work builds on these ideas, using empowerment as an exploration bonus that can efficiently leverage prior data.

Exploration from offline data. Offline RL learns policies purely from fixed datasets without further environment interaction. Offline-to-online methods pretrain a policy or value function on the offline data and then fine-tune it with online interaction [28] (Uchendu et al. 2023). Prior methods have shown that an online RL agent can be trained to explore much more efficiently in the environment by leveraging offline data [23] (Wilcoxson et al. 2025), but do not define a principled way to characterize the “exploration value” of a state. Our work instead learns empowerment offline as a principled exploration objective.

Preliminaries

We formulate the environment as a Markov Decision Process (MDP) defined by a state space $\mathcal{S}$, action space $\mathcal{A}$, transition dynamics $P(s' \mid s,a)$, and a discount factor $\gamma$. Let trajectories be denoted by the random variables $s_0,s_1,\ldots,s_T$. We are particularly interested in the distribution of discounted future states, denoted as $s^+ \triangleq s_{t+k}$ where $k \in\{0,1,2,\ldots\}$ is drawn with probability $(1-\gamma)\gamma^k$ [12]. Let $\pi(a \mid s,z) \in\Pi^z$ be a skill-conditioned policy, where $z \in\mathcal{Z} = \{1,\ldots,K\}$ is a latent skill drawn from a fixed, pre-determined set of $K$ skills and represented as a one-hot vector. In line with previous work [10], we define the discounted state occupancy measures as

Learning occupancy measures by temporal differences. Because $s^+$ is drawn from a geometric mixture over future time steps, the occupancy measures satisfy a Bellman-style recursion [10]:

This recursion can be used as a temporal-difference target to learn occupancy measures from off-policy transitions $(s,a,s')$, without rollouts of $\pi$ itself. By analogy with value functions, we write $Q^z(s^+\mid s,a)$ and $V^z(s^+\mid s)$ for learned models of $p^z(s^+\mid s,a)$ and $p^z(s^+\mid s)$, respectively.

Empowerment [19, 26] quantifies the maximal influence an agent can exert over its future outcomes from a given state, typically measured as the channel capacity between a control or intention and a future outcome. Prior work has offered many variations of the specific form. We use the definition from Gregor et al. [15], which defines the empowerment in state $s$ as the maximum possible mutual information between a uniformly-distributed random latent skill $z \in\mathcal{Z} = \{1,\ldots,K\}$ and the resulting future state $s^+$, conditioned on the current state:

We call this the potential or maximal empowerment, to distinguish it from the effective empowerment defined next. Intuitively, empowerment should be high when a policy can realize a diverse set of outcomes where each outcome corresponds to a distinct skill. For example, an agent stuck in the dead end of a maze has very low empowerment, as its actions have almost no outcome.

We take $\Pi^z$ to be the class of Gaussian policies with a learned mean $\mu_\theta(s,z)$ and a fixed standard deviation $\sigma_a$, clipped to the action bounds: $a=\operatorname{clip}(\mu_\theta(s,z)+\eta,-1,1)$ with $\eta\sim\mathcal{N}(0,\sigma_a^2 I)$. For a deterministic policy in a deterministic environment, each occupancy measure $p^z(\cdot\mid s)$ is a set of point masses on the states of a single trajectory, so it has no density over the continuous state space, and the mutual information theoretically saturates: $I^\pi(s^+;z\mid s)=\gamma\log K$ whenever no two skills’ trajectories share a state after $s$, no matter how close together they are. With fixed action noise, each skill’s occupancy measure is spread over a neighborhood of its trajectory instead, so the skills’ occupancy measures overlap and the mutual information reflects how far apart the skills’ futures are.

While potential empowerment represents the theoretical maximum control, evaluating it directly would require searching over all policies. Instead, we can define the effective empowerment as a lower bound induced by a specific, known policy $\pi$:

where the equality uses that $z$ is uniform; $\widetilde{\mathcal{E}}(s;\pi)\leq\mathcal{E}(s)$, with equality for the maximizing $\pi$.

Our method iteratively discovers a policy to maximize this effective empowerment, thereby approximating the potential empowerment $\mathcal{E}(s)$, using samples $(s,a,s')$ from an offline dataset $\mathcal{D}$.

Problem formulation. We assume access only to a static dataset $\mathcal{D}$ of transitions $(s,a,s')$ collected by an unknown behavior policy $\pi_\beta$, with no further environment interaction. Our goal is to use $\mathcal{D}$ to obtain an estimator $\widehat{\mathcal{E}}(s)$ of the potential empowerment $\mathcal{E}(s)$ (Eq. 5) at any state $s$.

An Algorithm for Offline Empowerment

In this section, we present our algorithm for estimating empowerment from offline data. The key insight is that the effective empowerment of a policy depends on that policy only through its discounted state occupancy measures (Eq. (6)), which can be learned off-policy using temporal differences (Eqs. (3)–(4)). Our method jointly learns occupancy models $Q_{\theta}^{z}(s^{+}\mid s,a)$ and $V_{\theta}^{z}(s^{+}\mid s)$ for a skill-conditioned policy by TD on offline transitions (§ 4.1), trains that policy $\pi(a\mid s,z)$ to maximize the effective empowerment computed from these models so that it approaches the potential empowerment $\mathcal{E}(s)$ (§ 4.2), and reads off an empowerment estimator $\hat{\mathcal{E}}(s)$ from these models (§ 4.3). Throughout training, we sample skills $z\sim p(z)$ uniformly from $\mathcal{Z}$.

Learning the Occupancy Measure

To evaluate effective empowerment offline (§ 3), we must estimate the distribution of future states without environment interactions. We achieve this by learning function approximators $Q_{\theta}$ and $V_{\theta}$ that model the discounted future state occupancy measures so that $Q_{\theta}^{z}(s^{+}\mid s,a)\approx p^{z}(s^{+}\mid s,a)$ and $V_{\theta}^{z}(s^{+}\mid s)\approx p^{z}(s^{+}\mid s)$. Each training example consists of a transition $(s,a,s')\sim\mathcal{D}$ together with a random state $s^{+}$ drawn uniformly from all states in $\mathcal{D}$, independently of the transition. Because the empowerment-maximizing policy differs from the behavior policy, these sampled futures are not technically samples from $p^{\pi^{z}}(\cdot\mid s,a)$, but the recursive backup allows us to still learn these quantities.

Bellman Backups. We train the occupancy networks by using the recursion of Eq. (3)–(4) as a TD target, with the next state $s'$ sampled from the dataset and the expectation over actions in Eq. (4) taken under the current policy $\pi(\cdot\mid s,z)$. A caveat is that the first term of Eq. (3), $(1-\gamma)1(s^{+}=s)$, is almost never observed, as a random state $s^{+}$ drawn from the dataset almost surely differs from the current state $s$. We resolve this by adding a second loss where $s^{+}=s$ instead of the randomly sampled offline data state, so that its target includes the $(1-\gamma)$ term. Equivalently, we sum two regression losses with equal weight,

with total objective $\mathcal{L}_{Q}=\mathcal{L}_{Q}^{\mathrm{boot}}+\mathcal{L}_{Q}^{\mathrm{self}}$, where $\ell$ is a clipped linear exponential regression loss used in prior works for stability (Parsian & Kirmani, [30]) (§ 4.4). The $V$-loss is

Note that $Q_{\bar{\theta}},V_{\bar{\theta}}$ denote target networks for Q and V. We update the networks using Polyak averaging (Lillicrap et al., [24]) with an update rate of $\tau=0.005$.

Latent Space Parameterization. Recall that computing empowerment (Eq. (6)) requires samples from $p^{z}(s^{+}\mid s)$. We now describe a parameterization of $V_{\theta}^{z}$ that makes such sampling tractable.

To tractably model these high-dimensional distributions, we project states and actions into a shared representation space of dimension $d$. We define non-linear encoders $\phi_{Q}(s,a,z)$, $\phi_{V}(s,z)$, and $\psi(s^{+})$, and parameterize the log-occupancies as scaled negative Euclidean distances:

This formulation allows us to sample directly from the discounted state-occupancy measure as an isotropic Gaussian in the latent space $p(\psi\mid\phi)=\mathcal{N}\left(\phi,\frac{d}{2}I\right)$, which is needed to compute the empowerment. Note that while we cannot sample $s^{+}$, we can still sample $\psi(s^{+})$ which is sufficient for computing $V^{z}(s^{+}\mid s)$ as we can plug the sampled latent into Eq. (10).

It is possible to use separate networks for both $\phi_{Q},\phi_{V}$, or alternatively, only learn $Q^{z}(s^{+}\mid s,a)$ and simply use $V^{z}(s^{+}\mid s)=\mathbb{E}_{a\sim\pi(\cdot\mid s,z)}[Q^{z}(s^{+}\mid s,a)]$. We find that using separate networks for Q and V leads to more stable training and qualitatively better estimates of the empowerment.

Policy Optimization via the Empowerment Gradient

Now that we can learn and sample from the state occupancies, we optimize a policy $\pi(a\mid s,z)$ to maximize its effective empowerment. Because empowerment depends only on the state, we need a way to describe how a skill’s action distribution influences it. We can define $\mathcal{M}^{z'}(s^+ \mid s,\pi,z)$ as the marginal modified discounted state occupancy measure, representing the distribution over $s^+$ if skill $z$ draws its action at $s$ from $\pi(\cdot\mid s,z)$ and all skills follow the current policy thereafter:

We can then define the $(\pi,z)$-modified empowerment,

which quantifies the local change in empowerment caused by changing the action distribution of skill $z$ at $s$. Increasing this encourages the skill-conditioned policy to take actions that make its future state occupancies under different skills distinguishable. Since we can easily sample from the marginal occupancy measure $V^{z'}$ via our Gaussian parameterization from above, we rewrite this as an importance-weighted expectation and estimate it by Monte Carlo sampling:

The complete policy objective is:

where $\alpha_{\mathrm{bc}}$ is a hyperparameter that enforces a behavioral constraint. The behavioral constraint is the log-likelihood of the dataset action under a unit-variance Gaussian centered at the policy mean. We train the $Q$, $V$ networks and the policy jointly, taking a gradient step on their combined losses at every iteration (Algorithm 1).

Estimating the Empowerment

Once training has converged, the learned policy $\pi$ approximately maximizes the effective empowerment, so its effective empowerment $\widetilde{\mathcal{E}}(s;\pi)$ serves as an estimate of the potential empowerment $\mathcal{E}(s)$. We compute it by substituting the learned occupancy model $V_\theta^z$ into Eq. (6) and approximating the expectation by sampling latents from the Gaussian parameterization of $V_\theta^z$ in Eq. (10):

where $\psi_{z,m}\sim\mathcal{N}(\phi_V(s,z),\frac{d}{2}I)$ are $M$ latent samples per skill and $V_\theta^z(\psi\mid s)$ denotes Eq. (10) evaluated at latent $\psi$. This is the estimator we evaluate in § 6.

Additional Details

Action noise. Every expectation over $a\sim\pi(\cdot\mid s,z)$, in $\mathcal{L}_V$ and in the policy objective (§ 4.2), is estimated with reparameterized samples $a=\operatorname{clip}(\mu_\theta(s,z)+\eta,-1,1)$. Since the policy’s actions are small perturbations of its mean, we also perturb the dataset actions with the same noise when training $Q_\theta$, $\tilde{a}=\operatorname{clip}(a+\eta,-1,1)$, so that $Q_\theta$ is trained on the same kind of slightly noisy actions at which it is later evaluated. Table 2 lists the value of $\sigma_a$ used for each dataset.

Clipped Linear Exponential Loss. Because $Q$ and $V$ operate in log-space, we regress in the exponentiated domain using a clipped linear-exponential loss for numerical stability, with the clipping threshold set to $t=5$ to avoid exploding gradients; its exact form is given in Section A.5.

Experimental Setup

Our experiments have two stages. We first train our empowerment estimator on offline OGBench datasets and compare it against offline adaptations of existing empowerment estimators (Sections 6.1 Algorithm 1: Offline Empowerment Estimation

Require: Offline dataset $\mathcal{D}$ of transitions $(s, a, s^\prime)$; number of skills $K$; discount $\gamma$; behavioral constraint coefficient $\alpha_{\mathrm{bc}}$; Polyak rate $\tau$

1. Initialize occupancy networks $Q_{\theta}$, $V_{\theta}$, target networks $\overline{Q}_{\overline{\theta}} \leftarrow Q_{\theta}$, $\overline{V}_{\overline{\theta}} \leftarrow V_{\theta}$, and policy $\pi(a \mid s, z)$

2. for each training step do

3. Sample a batch of transitions $(s, a, s^\prime) \sim\mathcal{D}$

4. Sample a random state $s^+ \sim\mathcal{D}$ and a skill $z \sim\operatorname{Unif}\{1,\ldots,K\}$ for each transition

5. Compute the occupancy losses $\mathcal{L}_Q = \mathcal{L}^{\mathrm{boot}}_Q + \mathcal{L}^{\mathrm{self}}_Q$ and $\mathcal{L}_V$ ▷ § 4.1

6. Sample latent futures $\psi^{z^\prime}_{z,m} \sim\mathcal{N}_{\phi_V}(s, z^\prime), \frac{d}{2} I$ for each $z^\prime\in\mathcal{Z}$, $m = 1,\ldots,M$

7. Sample actions $a_m \sim\pi(\cdot\mid s, z)$ with the reparameterization trick

8. Estimate $E_{\Delta}(s; \pi, z)$ with Eq. (13), averaging over the samples ▷ § 4.2

9. Compute the policy loss $\mathcal{L}_{\pi}$ from Eq. (14)

10. Take one gradient step on $\mathcal{L}_Q + \mathcal{L}_V + \mathcal{L}_{\pi}$

11. $\overline{\theta} \leftarrow\tau\theta+ (1-\tau)\overline{\theta}$ ▷ Polyak update of $\overline{Q}_{\overline{\theta}}$, $\overline{V}_{\overline{\theta}}$

12. return $\widehat{E}(s)$ computed from $V_{\theta}$ with Eq. (15) ▷ § 4.3

and 6.2). We then freeze the estimator and use it to guide online exploration on sparse-reward, goal-reaching tasks (§ 6.3). Since computing the true empowerment over all policies is intractable in continuous environments, we evaluate the estimator qualitatively, in environments whose structure dictates where the empowerment should be high (e.g., toward the center of a maze rather than in its dead ends). Hyperparameters are given in Table 1.

Hyperparameter	Value
Number of skills $K$	50
$Q/V$ network width / depth	512 / 3
Policy network width / depth	512 / 3
Value latent dimension $d$	256
Layer normalization	yes
Learning rate ($Q/V$ and policy)	$3 \times10^{-4}$
Batch size	1024
Discount $\gamma$	0.99
Target network update rate $\tau$	0.005
Clipped linexp threshold $t$	5
Action noise std $\sigma_a$	0.01
Behavioral constraint coefficient $\alpha_{\mathrm{bc}}$	0.001
Future-state samples per update	1
Training steps	$10^6$

Baselines

Our method is the first to learn empowerment in continuous environments from offline data, so there are no direct baselines to compare to; instead, we adapt four existing online empowerment estimators to the offline setting. We give the core idea behind each baseline below; full derivations, architectures, and hyperparameters are given in Section A.6.

Empowerment baselines. The baselines differ from each other, and from our method, along two orthogonal axes: which mutual information they measure and how they estimate it. Three measure the action empowerment $I(s^+; a \mid s)$, the mutual information between the action taken at $s$ in the dataset and the future state $s^+$ (Section A.6), and one measures the skill empowerment $I(s^+; z \mid s)$, as our method does; unlike our method, none of them computes it from learned occupancy measures. CF-Empowerment (contrastive features) estimates $I(s^+; a \mid s)$ with a binary NCE contrastive critic ([12]) using in-batch negatives. CF-FlowBC-Empowerment replaces those in-batch negatives with actions from a learned flow-matching behavior-cloned policy ([25]). DV-Empowerment is architecturally identical but optimizes the Donsker–Varadhan bound ([6]) instead of InfoNCE. DADS-Empowerment estimates the skill empowerment $I(s^+; z \mid s)$, matching our formalism, by jointly learning a discrete skill labeling and transition channel that maximize the Barber–Agakov bound ([34, 4]).

Online, errors in their learned models (e.g., dynamics or skill labelings) are corrected by fresh data, but offline, these methods suffer distribution shift and approximation errors that bias the estimate (§ 6.2).

Online Exploration with Offline Empowerment

We also use the empowerment estimator $\widehat{\mathcal{E}}(s)$ (Eq. (15)), trained on $\mathcal{D}$ and then frozen, to guide exploration when an agent later solves online, sparse-reward, goal-reaching tasks in the same environment.

Base agent. Our online agents are built on contrastive RL (CRL) ([12]), a goal-conditioned actor-critic. The critic $f(s,a,g) = \phi(s,a)^{\mathsf{T}}\psi(g)$ is trained with a contrastive loss to tell whether $g$ is a future state of $(s,a)$ on the same trajectory, and the policy $\pi(a \mid s,g)$ maximizes $f$ along with a SAC-style auto-tuned entropy term ([16, 17]). Every online agent is trained with Reinforcement Learning with Prior Data (RLPD) ([3]); each training batch is drawn half from the online replay buffer and half from $\mathcal{D}$.

Empowerment bonus. A naive way to use empowerment is to add it as an additional reward; however, this can have the undesired consequence of keeping the policy from leaving high-empowerment states. Instead, we would like the policy to visit high-empowerment states more often so it can explore around them, but not be penalized for leaving those high-empowerment states and continuing its trajectory into other parts of the state space. Consequently, we train a network $B_{\omega}(s,a)$ to predict the maximum empowerment reached anywhere along the trajectory containing the transition $(s,a)$. For every complete trajectory $\tau=(s_0,a_0,\ldots,s_T)$ in either the online or offline buffer, each of its transitions is regressed onto the highest empowerment reached along $\tau$:

Since the max is taken over both the future and the history, leaving a high-empowerment state is not penalized, while actions that lead to one are rewarded. The policy then maximizes

where $\beta$ is the bonus weight, and $\alpha_{\mathrm{ent}}$ is the automatically tuned entropy temperature.

Exploration baselines. We compare against standard CRL and two exploration methods. Random Network Distillation (RND) ([8]) rewards novel states, measured by how poorly a trained predictor matches a fixed, randomly initialized network; this reward is learned into a separate critic $Q_{\mathrm{RND}}(s,a)$, and the CRL policy maximizes $f(s,a,g)+\beta_{\mathrm{RND}}Q_{\mathrm{RND}}(s,a)$ along with the entropy term. SUPE ((Wilcoxson et al., 2025)) pretrains a latent skill model on $\mathcal{D}$, then trains a high-level SAC policy online that reselects a skill every $H$ steps; offline data is relabeled with a skill and mixed into training with RLPD. We anneal both the RND and empowerment weight linearly to zero 50% of the way through training, as we find this gives better results than holding constant.

Combining empowerment with other exploration methods. Because our bonus only adds a term to the policy objective, it can be combined with other exploration methods. We test the effect of combining our empowerment bonus with RND and also adding it to the SUPE high-level policy loss.

Evaluation. We evaluate on sparse-reward, goal-conditioned variants of OGBench environments in which the agent starts from a fixed state and must reach goals spread across the environment, as exploration is especially critical in these environments due to the sparsity of the reward signal: antmaze-medium (an ant starts in a corner of the medium maze and must reach other cells of the maze), pointmaze-teleport (a point must navigate a maze with random teleporters), antsoccer-arena (the ant must push a ball to different goals), and cube-single (a robot arm must move a cube to a different position on the ground). Section A.3.1 describes each task in detail. Each task is paired with one or more OGBench datasets for $\mathcal{D}$, which differ in data quality: navigate and play contain long, goal-directed trajectories, stitch contains only short trajectory segments that must be stitched together, noisy adds heavy action noise, and explore contains highly suboptimal, random-walk-like trajectories. The empowerment estimator for each task is trained on the same dataset $\mathcal{D}$ used by the online agent for RLPD. Implementation details and hyperparameters are given in Section A.3.

Experiments

We answer two questions: (1) Does our method recover meaningful empowerment estimates from offline data? (2) Can we use the offline empowerment estimator to improve online exploration?

Offline Empowerment Plots

Figure 1 shows the empowerment recovered by our offline estimator on three OGBench environments; in each case it matches what we intuitively expect. On antsoccer-arena-navigate-v0 it peaks immediately around the ball in both plots, where the ant can freely choose whether to approach, dribble, or move away from it. On antmaze-medium-navigate-v0 it peaks at the maze’s two corridor junctions and falls off in dead ends and corners. On pointmaze-teleport-navigate-v0 it collapses at each one-way teleporter entrance, where the agent has no control over its random next state.

Figure 2 shows the empowerment estimate along real dataset trajectories. In antsoccer (Figure 2a), empowerment rises as the ant approaches the ball, drops sharply the moment it kicks it, and stays low while the ball rolls out of reach, then climbs again as the ant walks back and dribbles, peaking just before the second kick with the ball at its feet. In cube (Figure 2b), empowerment rises as the gripper approaches and grasps the block. Once the arm carries the block up into the air, the empowerment falls: this is because the gripper must stay closed to retain control of the block, giving the agent fewer degrees of freedom to control its future state. Empowerment rises as the arm descends and releases the block; here empowerment is high because, as in antsoccer, the arm can choose to back off from the cube, or grab it and move it.

Qualitative Comparison to Baselines

Figure 3 shows $\hat{E}(s)$ recovered by the four empowerment baselines of § 5.1. None recovers a comparable map. DADS-Empowerment saturates at its ceiling of $\log K$ almost everywhere; CF-Empowerment learns a reasonable estimate, but is noisy and only captures the coarser trends; CF-FlowBC-Empowerment also fails. DV-Empowerment, whose bound is known to suffer from high variance (Song & Ermon, 2020), is virtually unbounded.

Online Exploration

Figure 4 compares our empowerment bonus against CRL, CRL + RND, and SUPE on every task and dataset (§ 5.2), and also shows the bonus combined with RND and SUPE. Across cube-single-play, cube-single-noisy, antmaze-medium-navigate, antmaze-medium-stitch, and pointmaze-teleport-stitch, adding the empowerment bonus improves over baseline CRL, enabling faster learning and higher performance. On the other environments, our method still performs on par with CRL and does not hurt performance. Our method can also be combined with existing exploration methods. Adding the empowerment to CRL + RND leads to improved performance on cube-single, antmaze-medium, and pointmaze-teleport over RND alone. Adding it to SUPE on cube and antmaze speeds up learning.

Ablations. Section A.2 ablates the action noise scale $\sigma_{a}$ and behavioral constraint coefficient $\alpha_{\mathrm{bc}}$ of our estimator, and the weight $\beta$ of the online empowerment bonus.

Conclusion

In this work we introduced a method for estimating empowerment purely from offline data, by jointly learning an empowerment-maximizing policy together with its discounted state-occupancy measure. Unlike prior methods, which rely on online rollouts of a skill-conditioned policy, our estimator trains on unlabeled offline trajectories and recovers an intuitive empowerment landscape across OGBench environments. Our results show that the estimator can be used to improve online exploration on sparse, goal-conditioned RL tasks. Beyond exploration, our estimator is general, and we hope future work will apply it to other problems and domains.

Appendix

Hyperparameters

This section shows hyperparameters for our offline empowerment estimator (§ 4). Hyperparameters for the four baselines are given alongside their algorithm descriptions in Section A.6.

Offline Empowerment (Ours). Table 1 lists hyperparameters for the offline empowerment estimator and skill-conditioned policy of § 4.

Environment	Dataset	$\sigma_a$
antmaze-medium	navigate	0.01
antmaze-medium	stitch	0.01
antmaze-medium	explore	0.01
antsoccer-arena	navigate	0.01
antsoccer-arena	stitch	0.01
pointmaze-teleport	navigate	0.01
pointmaze-teleport	stitch	0.01
cube-single	play	0.01
cube-single	noisy	0.01

Ablations

This section ablates three hyperparameters of our method: the policy’s action noise scale $\sigma_a$ (§ 3 and § 4.4), the behavioral constraint coefficient $\alpha_{\mathrm{bc}}$ (§ 4.2), and the weight $\beta$ of the online empowerment bonus (§ 5.2). For $\sigma_a$ and $\alpha_{\mathrm{bc}}$, we sweep the hyperparameter, holding everything else at the values in Table 1, and track the validation empowerment estimate during training.

Action Noise Scale

Figure 5 sweeps $\sigma_a$ over three orders of magnitude, $\sigma_a \in\{10^{-4}, 10^{-3}, 10^{-2}, 10^{-1}\}$, on antmaze-medium-navigate-v0. The converged empowerment estimate is highly robust to the action noise scale.

Behavioral Constraint Coefficient

Figure 6 sweeps the behavioral constraint coefficient $\alpha_{\mathrm{bc}} \in\{3 \times10^{-4}, 10^{-3}, 3 \times10^{-3}, 10^{-2}\}$ on antsoccer-arena-stitch-v0. Because $\alpha_{\mathrm{bc}}$ directly weights how strongly the policy is pulled toward the offline behavior policy, a value that is too large prevents the policy from ever departing from the dataset’s own behavior far enough to realize the higher-empowerment behaviors. Conversely, $\alpha_{\mathrm{bc}}$ cannot be too low as it must regularize the skill-conditioned policy to the offline data.

Qualitative failure mode. Figure 7 shows the empowerment map for two checkpoints on antsoccer-arena-stitch-v0 that differ only in $\alpha_{\mathrm{bc}}$. At $\alpha_{\mathrm{bc}} = 10^{-3}$, empowerment peaks tightly around the ball, matching the desired pattern in Figure 1a. At $\alpha_{\mathrm{bc}} = 10^{-2}$, the map instead shows a distinctive hole exactly at the ball. This is the collapse from Figure 6 made spatially concrete: on the stitch dataset, the offline trajectories that pass near the ball are dominated by a single behavior—dribbling it—so a heavily BC-regularized policy is pulled toward dribbling near the ball regardless of skill, and never discovers the alternative high-empowerment behaviors (e.g. deliberately moving away from the ball) that would make the skill distinguishable. Notably, this failure mode is specific to the stitch dataset; since the navigate dataset contains trajectories of the ant doing many behaviors near the ball, a much wider range of BC values works well. On low-coverage data such as stitch, $\alpha_{\mathrm{bc}}$ must be kept low enough so that the policy can generalize beyond the offline data and learn empowerment-maximizing skills. In general, we find that there is around a single order of magnitude range in which the BC coefficient should lie to give good empowerment estimates, for all environments.

Empowerment Bonus Weight

Fig. 8 varies the weight $\beta$ of the empowerment bonus in the policy objective of Eq. (17), comparing with CRL ($\beta= 0$). We sweep $\beta\in\{3, 10, 30\}$ on cube-single (play) and antmaze-medium (navigate). On cube-single, performance is non-monotonic in $\beta$: $\beta= 10$ performs best, while $\beta= 3$ improves only slightly over CRL. On antmaze-medium, every $\beta$ improves over CRL. In general, like other exploration methods, $\beta$ must be chosen to encourage the policy to visit high-empowerment states without dominating the policy objective.

Online Exploration Details

This section gives implementation details for the online experiments of § 5.2.

Online Environments

All online tasks are goal-conditioned variants of OGBench environments (Park et al., [29]) with a single fixed start state and a fixed set of goals. At the start of every episode, one goal is drawn uniformly from the set and given to the agent, and neither the start nor the goal is perturbed by noise. The reward is sparse: the agent receives no signal until it reaches the goal, and an evaluation episode counts as a success if the agent reaches the goal within the episode horizon. Table 3 summarizes the tasks, and the offline datasets used for RLPD and for pretraining the empowerment estimator. In Figure 4, shaded regions show $\pm1$ standard error of the mean over seeds.

Task	Goals	Horizon	Offline datasets
antmaze-medium	25	1000	navigate, stitch, explore
pointmaze-teleport	41	1000	navigate, stitch
antsoccer-arena	8	500	navigate, stitch
cube-single	30	200	play, noisy

• antmaze-medium: The ant starts in the bottom-left corner of OGBench’s medium maze, and the goal is the center of any of the 25 other free cells of the maze, so most goals require navigating several corridors and turns.

• pointmaze-teleport: A point mass starts in the bottom-left corner of the teleport maze, and the goal is the center of any of 41 other free cells. The maze contains teleporters that send the agent to a random exit; we exclude the two teleporter entrances, where the agent cannot remain inside the goal region, and one walled-off pocket that can only be entered by a lucky teleport and cannot be left.

• antsoccer-arena: The ant starts in the open arena about three units behind the ball, so it must first walk to the ball and then push it to one of eight goals, each one to three units beyond the ball’s start on the side away from the ant. The episode horizon is 500 steps.

• cube-single: A single cube starts at a fixed position on the table, and the goal is one of 30 positions on a $5 \times6$ grid covering the table. The closest goals are just outside the success radius of the start, so every task requires picking up or pushing the cube. The episode horizon is 200 steps.

Agents

Training loop. All online agents interact with a single environment. After an initial warm-up period, the agent collects a fixed number of environment steps and then performs a fixed number of gradient updates per environment step, on batches drawn from the online replay buffer and the offline dataset $\mathcal{D}$ in equal proportion (RLPD). For both sources, CRL goals are future states of the same trajectory, sampled at a geometrically distributed offset.

Empowerment bonus. The estimator $\hat{\mathcal{E}}$ is the pretrained, frozen model of § 4.3, trained on the same dataset $\mathcal{D}$ that RLPD uses. We evaluate $\hat{\mathcal{E}}(s)$ once for every state of $\mathcal{D}$ before training, and once for every new online state in a single batched call before each round of updates. The regression target in Eq. (16) is shifted by the mean of $\hat{\mathcal{E}}$ over $\mathcal{D}$; this shifts $B_{\omega}$ by a constant and leaves the policy gradient unchanged, but keeps the regression well conditioned. $B_{\omega}$ has two heads, both trained with Eq. (16), and the policy uses their minimum. Transitions of an online episode that has not yet ended are excluded from the regression until the episode ends and its target is known; they are used by the CRL critic and the policy as usual.

RND. The target and predictor networks share one architecture and take the next state $s'$ as input, normalized by a running mean and standard deviation of online states and clipped to $[-5,5]$. The intrinsic reward is the squared error between the predictor and target outputs, divided by a running standard deviation of the discounted intrinsic return. $Q_{\mathrm{RND}}$ is trained with a twin-critic Bellman backup ([14]) that does not stop at episode ends, as in [8]. The predictor, the running statistics, $Q_{\mathrm{RND}}$, and the policy’s RND bonus all use online transitions only.

CRL + RND + Empowerment. The two bonuses are trained independently exactly as above, with their own weights $\beta$ and $\beta_{\mathrm{RND}}$.

SUPE. We follow the implementation of Wilcoxson et al. (2025). The skill model is an OPAL-style VAE (Ajay et al. [1]) with a continuous latent, trained on $\mathcal{D}$ over segments of $H = 4$ steps (8 for soccer) and then frozen. The high-level policy is a tanh-Gaussian SAC actor over the latent, with an ensemble of 10 layer-normalized (Ba et al. [2]) critics. Each offline segment inside a single trajectory is labeled with the posterior mean of its latent. Following SUPE’s treatment of sparse $-1/0$ tasks, at every update its reward is set to the dataset minimum plus an RND bonus, and its termination is set to the prediction of a reward model trained on online data.

SUPE + Empowerment. $B_{\omega}(s,u)$ has the same architecture as for CRL, takes the high-level state and skill $u$ as input, and is regressed onto the same targets as in Eq. (16). The high-level actor loss becomes $\alpha_{\mathrm{ent}}\log\pi(u \mid s,g) - \overline{Q}(s,u,g) - \beta B_{\omega}(s,u)$, where $\overline{Q}$ is the mean over the critic ensemble.

Hyperparameter selection. For each task and dataset, we perform a sweep over the weight of each exploration bonus and report the best-performing setting; $\beta\in\{3, 10, 30\}$ for the empowerment bonus and $\beta_{\mathrm{RND}} \in\{0.3, 1, 3\}$ for RND. For SUPE + Empowerment, we sweep $\beta\in\{10, 30\}$. For CRL + RND + Empowerment, we evaluate two configurations: the best weights from the individual sweeps and both weights halved. Table 6 lists the weights used for each task; all other hyperparameters are shared across tasks (Table 4 and Table 5).

Hyperparameter	Value
Actor network width / depth	512 / 3
Critic network width / depth	512 / 3
Critic latent dimension	512
Layer normalization	yes
Learning rate	$3 \times10^{-4}$
Batch size	1024
Offline fraction of each batch (RLPD)	0.5
Discount $\gamma$	0.99
Target entropy	$-\frac{1}{2}\dim(\mathcal{A})$
Environment steps per update round	50
Updates per environment step	1
Warm-up environment steps	1000
Replay buffer size	$10^6$

Hyperparameter	Value
Empowerment bonus (ours)
$B_\omega$ network width / depth	512 / 3 (two heads)
Bonus weight $\beta$	per task (Table 6)
Future-state samples per skill for $\hat{\mathcal{E}}$	64
RND
Target / predictor network width / depth	512 / 2
Embedding dimension	512
Predictor learning rate	$10^{-4}$
$Q_{\mathrm{RND}}$ network width / depth	512 / 3 (two heads)
$Q_{\mathrm{RND}}$ discount	0.99
$Q_{\mathrm{RND}}$ target update rate $\tau$	0.005
Bonus weight $\beta_{\mathrm{RND}}$	per task (Table 6)
SUPE
Skill horizon $H$	4 (8 on antsoccer-arena)
Skill latent dimension	8
OPAL KL coefficient	0.1 (0.2 on cube-single)
Number of critics / critics in target minimum	10 / 1 (10 / 2 on cube-single)
Discount	0.995
RND coefficient	8
Warm-up environment steps	5000

Task	Dataset	CRL + Emp. $\beta$	CRL + RND $\beta_{\mathrm{RND}}$	CRL + RND + Emp. $\beta$ / $\beta_{\mathrm{RND}}$	SUPE + Emp. $\beta$
cube-single	play	10	1	10 / 1	30
cube-single	noisy	30	1	30 / 1	10
antmaze-medium	navigate	30	3	30 / 3	10
antmaze-medium	stitch	10	3	10 / 3	10
antmaze-medium	explore	10	3	10 / 3	10
pointmaze-teleport	navigate	30	1	15 / 0.5	30
pointmaze-teleport	stitch	10	0.3	10 / 0.3	30
antsoccer-arena	navigate	10	1	10 / 1	10
antsoccer-arena	stitch	10	1	10 / 1	10

Compute Resources

All offline pretraining runs were trained on an internal cluster of NVIDIA A5000 and A6000 GPUs, with a single GPU per run.

Clipped Linear Exponential Loss

Because $Q$ and $V$ operate in log-space (§ 4), we regress in the exponentiated domain for numerical stability. For a target $y$ and prediction $\hat{y}$, we regress $e^{-\hat{y}}$ to $\kappa e^{-y}$ (for a target scale $\kappa$) using

The second branch is simply a linearization of the loss about the point $\hat{y}=y+t$, applied whenever $\hat{y}-y$ exceeds the threshold $t=5$, which prevents the exponential term from producing exploding gradients.

Baseline Algorithm Details

This section gives full derivations, network architectures, and hyperparameters for the four baselines summarized in § 5.1. All four baselines are trained on triples $(s,a,s^+) \sim\mathcal{D}$, where $s^+$ is a future state of $s$ on the same dataset trajectory, drawn at a geometrically distributed offset $k \ge1$, with probability $(1-\gamma)\gamma^{k-1}$. This differs from our method, which draws $s^+$ uniformly from $\mathcal{D}$ and recovers the dependence on $(s,a)$ through TD backups. With the exception of DADS-Empowerment (Section A.6.4), the baselines estimate the action empowerment

that is, the mutual information between the primitive action $a \sim\pi_\beta(\cdot\mid s)$ actually taken at $s$ under the behavior policy $\pi_\beta$ and the resulting future state $s^+$, in place of the mutual information between a skill $z$ and $s^+$ used by our method.

CF-Empowerment

Overview. CF-Empowerment estimates the action empowerment $I(s^+;a\mid s)$ by first fitting two contrastive critics that share a future-state encoder, then distilling the resulting pointwise mutual-information estimator into a small feedforward network that can be evaluated at any state without re-running the contrastive critics.

Architecture. Three MLP trunks map into a shared latent space of dimension $d$: an action-conditioned encoder $\phi(s,a)$, a state-only encoder $\phi'(s)$, and a future-state encoder $\psi(s^+)$, all mapping into $\mathbb{R}^d$. A separate empowerment network $E_\theta(s) \in\mathbb{R}$ is a feedforward MLP that only depends on the current state. There is no policy network.

Training objective. Given a minibatch of $B$ triples $\{(s_i,a_i,s_i^+)\}_{i=1}^B$, define the bilinear logits

The critics are trained with a binary NCE objective ([12]), treating $s_i^+$ as the positive for $(s_i,a_i)$ (resp. $s_i$) and every other $s_j^+$ in the batch as a negative:

At the optimum, $\ell^a(s,a,s^+) \approx\log p(s^+\mid s,a)-\log p(s^+) + c$ and $\ell^s(s,s^+) \approx\log p(s^+\mid s)-\log p(s^+) + c$ for the same constant $c$, so $T(s,a,s^+) = \ell^a(s,a,s^+) - \ell^s(s,s^+)$ estimates the conditional mutual-information integrand $\log p(s^+\mid s,a)-\log p(s^+\mid s)$. After a fixed number of pretraining steps, the critics’ parameters are frozen and the empowerment network is trained by regression onto $T$:

At convergence $E_\theta(s) \to\mathbb{E}_{(a,s^+)\sim p(a,s^+|s)}[T(s,a,s^+)] = I(s^+;a\mid s)$, i.e., the network converges to an estimate of the action empowerment at $s$.

Offline adaptation. In the online setting, the negatives for the contrastive critics would be drawn by rolling out many actions from the same state; here they are instead the other $(s_j,a_j,s_j^+)$ triples already present in the offline minibatch (in-batch negatives), and $a$ ranges only over the fixed offline behavior policy rather than an arbitrary exploration policy.

Hyperparameters. Learning rate $10^{-3}$; batch size 512; critic hidden dims $(256,256)$ with layer normalization; latent dimension $d=256$; empowerment network hidden dims $(256,256)$; contrastive pretraining for $5 \times10^5$ steps before switching to distillation.

CF-FlowBC-Empowerment {#sec: A.6.2}

Overview. CF-FlowBC-Empowerment also estimates the action empowerment $I(s^+;a\mid s)$, but replaces the in-batch negatives of CF-Empowerment with negatives drawn from a learned behavior-cloned policy, and replaces the bilinear critic with a single scalar critic $T(s^+,a,s)$ trained with an explicit InfoNCE lower bound.

Architecture. Two flow-matching models ([25]) are trained on the offline data: a behavior-cloned policy $\pi_{\mathrm{bc}}(a\mid s)$ and a dynamics flow $f_{\mathrm{dyn}}(s^+\mid s,a)$ modeling the discounted future state. Both share the same velocity-field MLP architecture $v_\theta(x_t,t,\operatorname{cond})$, conditioned on $s$ (for $\pi_{\mathrm{bc}}$) or on $(s,a)$ (for $f_{\mathrm{dyn}}$); samples are drawn by Euler-integrating the learned velocity field for 8 steps from Gaussian noise, with sampled actions clipped to the action bounds. A separate scalar critic $T(s^+,a,s)$, an MLP over the concatenation of its three inputs, provides the mutual-information estimate.

Training objective. Both flow models are trained with the standard conditional flow-matching loss,

instantiated with $(\operatorname{cond},x)=(s,a)$ for $\pi_{\mathrm{bc}}$ and $(\operatorname{cond},x)=((s,a),s^+)$ for $f_{\mathrm{dyn}}$. Given a positive triple $(s,a,s^+)$, $N$ negative actions $a_1,\ldots,a_N \overset{\mathrm{iid}}{\sim} \pi_{\mathrm{bc}}(\cdot\mid s)$ are sampled and the critic is trained to maximize the categorical InfoNCE bound

The reported estimate $\hat{I}(s^+;a\mid s)=\log(N+1)-\mathcal{L}_{\mathrm{critic}}$ is therefore capped at $\log(N+1)$ nats, a known finite-sample ceiling of the InfoNCE bound ([31]). To evaluate the empowerment at an arbitrary state $s$ (rather than at a specific offline triple), $M$ joint samples $a\sim\pi_{\mathrm{bc}}(\cdot\mid s)$, $s^+\sim f_{\mathrm{dyn}}(\cdot\mid s,a)$ are drawn and scored against fresh $\pi_{\mathrm{bc}}$ negatives via the same bound and averaged. Because this substitutes $\pi_{\mathrm{bc}}$ and $f_{\mathrm{dyn}}$ for the true behavior policy and its discounted future-state distribution, the resulting estimate is the mutual information under the learned models rather than the true environment, and errors in either model bias it.

Offline adaptation. $\pi_{\mathrm{bc}}$ substitutes for the ability to roll out arbitrary actions from a live policy, and $f_{\mathrm{dyn}}$ substitutes for rolling out $\pi$ over a discounted horizon; both are trained on the offline dataset with their own flow-matching losses, and no gradient from the critic flows into them.

Hyperparameters. Learning rate $3 \times10^{-4}$; batch size 256; flow and critic hidden dims $(512,512,512)$ with layer normalization; 8 Euler integration steps; $N=7$ negatives; $10^5$ steps of $\pi_{\mathrm{bc}}/f_{\mathrm{dyn}}$-only pretraining before the critic begins training.

DV-Empowerment {#sec: A.6.3}

Overview. DV-Empowerment is architecturally identical to CF-FlowBC-Empowerment—the same $\pi_{\mathrm{bc}}$, $f_{\mathrm{dyn}}$, and scalar critic $T(s^+,a,s)$—but replaces the categorical InfoNCE bound with the Donsker– Varadhan (DV) lower bound [6], and replaces the per-element action negatives with full (action, future-state) negative pairs drawn from an explicit product-of-marginals construction.

Training objective. Given a positive triple $(s, a, s^{+})$, $N$ negative actions $a_1,\ldots,a_N \mathrel{\overset{\mathrm{iid}}{\sim}} \pi_{\mathrm{bc}}(\cdot\mid s)$ are sampled, each paired with its own simulated future state $s^{+}_{i} \sim f_{\mathrm{dyn}}(\cdot\mid s, a_i)$. To obtain valid samples from the product of marginals $\pi_{\mathrm{bc}}(a \mid s) \times f_{\mathrm{dyn}}(s^{+} \mid s, a)$ rather than the joint, the future states are re-paired with a cyclic derangement $\sigma$ (e.g. $\sigma(i) = (i \bmod N) + 1$), so that $(a_i, s^{+}_{\sigma(i)})$ never uses the future state that $a_i$ itself generated. The DV objective is then

with reported estimate $\hat{I}(s^{+};a \mid s) = \mathbb{E}[J(s,a,s^{+})]$. As with CF-FlowBC-Empowerment, evaluating $\hat{\mathcal{E}}(s)$ at an arbitrary state draws 512 fresh joint samples from $\pi_{\mathrm{bc}}$ and $f_{\mathrm{dyn}}$ and scores them against 512 negatives built with the same derangement-based construction.

Offline adaptation. Identical to CF-FlowBC-Empowerment: $\pi_{\mathrm{bc}}$ and $f_{\mathrm{dyn}}$, both fit on the offline dataset, substitute for on-policy action and rollout sampling.

Hyperparameters. Identical to CF-FlowBC-Empowerment: learning rate $3 \times10^{-4}$; batch size 256; flow and critic hidden dims (512, 512, 512) with layer normalization; 8 Euler integration steps; $N = 7$ negatives; $10^{5}$ steps of pretraining.

DADS-Empowerment

Overview. DADS-Empowerment is the only baseline that estimates the skill empowerment $I(s^{+}; z \mid s)$ for a small, fixed-size discrete skill alphabet, matching the formalism used by our method. Rather than fixing a skill-labeling scheme first, as in DIAYN/DADS-style skill discovery [34], and only then measuring the resulting mutual information, DADS-Empowerment learns the skill assignment itself so as to directly maximize the Barber–Agakov (BA) variational bound [4] on offline $(s,s^{+})$ pairs.

Architecture. Three networks operate on a fixed number $K$ of discrete, one-hot skills. An encoder $\rho(z \mid s,s^{+})$ is an MLP that maps the concatenation of $s$ and the observed transition delta $s^{+} - s$ to a softmax distribution over the $K$ skills; conditioning only on the endpoint pair (rather than any intermediate trajectory) ensures $z$ cannot encode information that never manifests in $s^{+}$ itself. A transition channel $q(s^{+} \mid s,z)$ is a diagonal-Gaussian mixture with 8 components conditioned on $s$ and a one-hot $z$. A skill marginal $m(z \mid s)$ is an MLP over $s$ producing a softmax over the $K$ skills. There is no action-conditioned network of any kind, and no policy: the primitive action never appears in this baseline’s training objective.

Training objective. For every offline pair $(s,s^{+})$ and every skill $z \in\{1,\ldots,K\}$ (the expectation over $z$ is computed exactly, not by Monte Carlo, since $K$ is small), define the BA bracket

where the log-sum-exp over $z'$ uses a skill marginal $m$. The encoder loss maximizes the bracket under its own soft skill assignment:

and the reported estimate is $\hat{I}(s^{+};z \mid s) = -\mathcal{L}_{\mathrm{bracket}}$. The skill marginal $m(z \mid s)$ is separately fit by cross-entropy toward the current (gradient-stopped) soft assignment,

so that $m$ tracks $\rho$’s marginal usage of each skill without corrupting $\rho$’s own gradient. A skill-usage regularizer, weighted by a coefficient $\lambda_{\mathrm{kl}}$,

where $\mathcal{U}_K$ is the uniform distribution over the $K$ skills. This term penalizes deviation of the average skill usage from uniform, discouraging the assignment from collapsing onto a small subset of skills. The total loss is $\mathcal{L}_{\mathrm{bracket}} + \mathcal{L}_{\mathrm{marginal}} + \lambda_{\mathrm{kl}}\mathcal{L}_{\mathrm{usage}}$.

Offline adaptation. As with the other three baselines, $(s, s^{+})$ pairs are drawn from the offline dataset; no skill-conditioned rollouts, dynamics-model rollouts, or EM alternation are used.

Hyperparameters. Learning rate $3 \times10^{-4}$; batch size 1024; mixture-density network with 8 components; encoder, channel, and marginal networks all use hidden dims (512, 512, 512) with layer normalization. We use $K = 15$ skills and $\lambda_{\mathrm{kl}} = 1$, including for the qualitative maps in § 6.2.