Offline empowerment: Learn where to explore from old data

Say you have a pile of old data with no rewards or skill labels, and a sparse-reward task to learn online. How can that data tell you where to explore? From the data alone, this paper learns skills that reach futures as different as possible, which gives an estimate of empowerment: roughly how many distinct futures you can reach from a state. A skill here is just a number z from 1 to K (50 in the paper) that goes into a single policy with the state, and nobody decides what each number means. Training does, by pushing different skills toward different futures, and empowerment is how much knowing z then tells you about where you'll go. In the ant maze, it's high at corridor junctions and low in corners. Online, it becomes an exploration bonus. On antmaze-medium with navigate data, it lifts contrastive RL from 0.29 to 0.43 success and also helps RND and SUPE. Across nine task and dataset pairs, it helps contrastive RL on five and ties on four.

From a state $s$, a policy $\pi$ can follow any of $K$ skills $z$, and you pick one uniformly at random. Call $s^+$ the discounted future state: a state some random number of steps ahead, about a hundred on average. $p^z(s^+ \mid s)$ is where skill $z$'s futures land, and $\bar p$ is their average over skills. Empowerment $E(s)$ is how much knowing the skill tells you about the future, their mutual information $I$, under the best policy. With uniform skills, that's the average KL divergence $D_{\mathrm{KL}}$, a measure of how different two distributions are, between each skill's futures and the average: $$\begin{aligned} E(s) &= \max_{\pi}\, I(s^+; z \mid s) \\ &= \max_{\pi}\, \frac{1}{K}\sum_{z=1}^{K} D_{\mathrm{KL}}\big(p^z(\cdot \mid s)\,\big\|\,\bar p(\cdot \mid s)\big) \end{aligned}$$ It's near zero when no policy can send the skills to different futures. You can't run skills on logged transitions $(s, a, s')$ (state, action, next state), so the method learns $V^z(s^+ \mid s)$, a model of $p^z(s^+ \mid s)$, with temporal differences like a value function, while it trains $\pi(a \mid s, z)$ to push the skills' futures apart. Plugging $V^z$ of the trained policy into the KL sum gives an estimate $\hat E(s)$, which then stays frozen. Online, contrastive RL's critic $Q(s, a, g)$ scores how well action $a$ from $s$ heads toward goal $g$. A bonus network $B_\omega$ with weights $\omega$ learns $B_\omega(s_t, a_t) \approx \max_{0 \le k \le T} \hat E(s_k)$, so every step $t$ of an episode that runs from step $0$ to $T$ gets the same target: the highest estimate at any step $k$ of that episode. The actor maximizes $Q(s, a, g) + \beta B_\omega(s, a)$ plus the usual entropy term, where $\beta$ weights the bonus. A plain empowerment reward can keep the agent parked at junctions. With the max, reaching a junction pays and leaving it to explore costs nothing.

← All writing