Learning to Perceive the World Through Control: Empowerment-Based Representation Learning
Mahsa Bastankhah, Sophie Broderick, Benjamin Eysenbach
cs.LG
2026-05-29
MISL representations vary at least ~1,000× more on controllable state than noise in maze, Ant, and pixel kitchen; frozen, they beat raw noisy inputs on reward.
In control, most of an observation is ballast. A maze agent needs coordinates; a robot needs joints and contact. Wall color, a drifting noise channel, and a pixel patch that never enters the dynamics do not change what can be done. Bisimulation merges states by reward and transitions, so it needs a reward, and any control-relevant factor that reward ignores can vanish with the noise. Inverse-dynamics features on a fixed dataset are invariant to uncontrollable factors only when the policies that collected the data already ignored them. Otherwise the dataset's bias becomes the representation.
The question is whether interaction alone, with no reward and no curated trajectories, can keep control-relevant information and drop the rest.
The objective is empowerment, used interchangeably with mutual information skill learning (MISL). A latent skill z indexes a family of policies. Empowerment at a start state is the mutual information between z and a future state drawn from that skill's discounted occupancy. More distinguishable futures mean higher empowerment. The practical objective is a variational lower bound: a discriminator sees the start and the reached state and infers which skill ran.
The discriminator is factored as q(z | ϕ(s0), ψ(s+)). The forward map ϕ clusters states by the futures they can influence. The backward map ψ clusters states by how they are reached. They do not alias the same pairs. Two cells at one maze location with different wall colors share dynamics, so they share ϕ. A locked door and a door that is open with probability 0.5 can be reached by the same action sequences from any start, so those cells share ψ. The open door affords much more future control, so ϕ differs. Several MISL implementations tie the two encoders together. The objective does not ask for that.
States split into y, w, and e. Actions move y. w is not moved by actions but changes how y moves: friction, or whether a door is open. e moves nothing and is moved by nothing else. Control-relevant information is x = (y, w). At least one empowerment-maximizing skill distribution ignores e. Future e carries no information about the skill, and acting on e does not enlarge the reachable distributions of future x. Any policy that reads e can be matched in discounted occupancy by a policy that reads only x. The matching policy may be nonstationary, and nonstationary policies do not expand that set. A representation that keeps only what predicts the skill therefore drops e in both ϕ and ψ.
That is an existence result. Other optima may still read e. Policy minimality selects among them: given the same occupancy, take the policy with the smallest state-action mutual information. A separate result covers action-interface relabeling. If one binding uses A to jump and another uses B, while the transitions match, MISL takes the corresponding action and the forward representations coincide. Successor representations do not. They follow the behavior policy that collected the data, so an interface-dependent collector yields an interface-dependent embedding.
Forward and backward features together are still not enough in general. Proposition 5.7 gives MDPs where the pair misses part of x. Proposition 5.8 adds a sufficient condition: x-transitions are deterministic, and any two x-states are both reachable in n steps from some common start. Then the backward representation of n-step empowerment separates every distinct x. For rewards that depend only on x, that representation is a sufficient statistic and matches the most general bisimulation that keeps all control-relevant features. Bisimulation on one reward also ignores e, and it can delete control-relevant features that reward never touches.
The implementation is METRA. Its original objective is a Wasserstein distance; later work shows equivalence to mutual information. Each observation is concatenated with Gaussian noise that is correlated in time and does not affect the dynamics. Sensitivity is the trace of the representation covariance while one coordinate moves and the others stay fixed.
| Environment | State size | Noise size | Representation variance |
| Point Maze | 2 | 5 | controllable dims ≥ 10^3 × noise dims |
| Ant | 29 | 10 | same order |
| Pixel Kitchen | 64×64×3 | 10×10×3 | same order |
Across the three domains, variance along controllable dimensions is approximately at least a thousand times the variance along irrelevant noise.
Frozen features are then trained with the environment's dense reward. Point Maze and Kitchen use SAC. Ant uses PPO. The baseline receives the raw noisy state. All three learning curves sit substantially above that baseline. The paper does not report a numeric return gap.
On noisy Point Maze, a linear probe reads the coordinates (x, y) out of the representation, and R² on both axes approaches 1. Training noise is 5-dimensional Gaussian, uncorrelated in time. Test noise is temporally correlated or correlated across dimensions, with the temporal coefficient plotted out to 0.99. The embedding stays nearly flat in those noise directions.
A tabular MDP checks the aliasing claim. States b and c have average posterior ℓ1 distance 2 backward and 0 forward. States e and f have forward distance 1 and backward distance 0. Forcing the backward posteriors of b and c to match drops empowerment at state A from 0.69 to 0.
The runs add no information bottleneck. The Point Maze policy is a 2-layer MLP of width 1024 with a 2-dimensional skill. Scaling the policy to 5 layers of width 1536, the discriminator to 8 layers of width 2048, and the skill to 16 dimensions leaves the noise invariance in place at convergence.
Empowerment has mostly been a recipe for diverse skills and exploration. The representation claim is sharper. If the agent interacts in order to make futures distinguishable, the encoders organize around control, and the forward and backward views are not interchangeable. Tying them together takes empowerment at that tabular state from 0.69 to 0.
The usable recipe is reward-free pretraining: run METRA, freeze the encoder, then train SAC or PPO. Inverse dynamics and successor representations are noise-invariant only if the data already came from a noise-invariant policy. Empowerment's objective is what produces that policy, so a cleaned dataset is not a prerequisite.
The main control plots compare only against raw observations. In the appendix, bisimulation tracks temporally correlated noise because it relies on a latent forward predictor, and AC-state inherits whatever the collection policy attended to. Neither comparison reports return against MISL. There is no basis yet for a ranking among representation methods.
The proofs are for fully observable MDPs. Partial observability is open. So is a question the paper states directly: after e is gone, it is unclear whether the representation recovers the causal structure among control-relevant variables, and unclear what that structure would be worth.
Keeping all of x is proved only when x-transitions are deterministic and any two x-states are jointly reachable in n steps. Without those assumptions there is a counterexample. Stochastic pixel control and contact-rich dynamics sit outside the sufficient condition. An R² near 1 shows that Point Maze coordinates survive in the embedding, not the general claim.
Noise invariance in the theorem uses policy minimality to rule out other optima, and it assumes a discriminator that reaches the Bayes posterior. METRA implements neither constraint. Invariance after widening the networks is evidence about this implementation, not a replacement for the assumption. The distractors are concatenated Gaussian channels or pixel blocks, not background video, lighting, or other agents. Downstream gains have no reported return delta and no learning curves against bisimulation or AC-state. Interface invariance is checked only in an appendix, with no numbers in the main text.