Soft Forward-Backward Representations for Zero-shot Reinforcement Learning with General Utilities
This work addresses the challenge of handling arbitrary differentiable objectives in RL, extending zero-shot capabilities beyond additive rewards for broader applications, though it builds incrementally on existing forward-backward methods.
The paper tackled the problem of zero-shot reinforcement learning with general utilities, which includes tasks like distribution matching and pure exploration, by introducing a soft forward-backward algorithm that recovers stochastic policies from offline data and optimizes utilities directly at test-time, demonstrating favorable properties across experiments.
Recent advancements in zero-shot reinforcement learning (RL) have facilitated the extraction of diverse behaviors from unlabeled, offline data sources. In particular, forward-backward algorithms (FB) can retrieve a family of policies that can approximately solve any standard RL problem (with additive rewards, linear in the occupancy measure), given sufficient capacity. While retaining zero-shot properties, we tackle the greater problem class of RL with general utilities, in which the objective is an arbitrary differentiable function of the occupancy measure. This setting is strictly more expressive, capturing tasks such as distribution matching or pure exploration, which may not be reduced to additive rewards. We show that this additional complexity can be captured by a novel, maximum entropy (soft) variant of the forward-backward algorithm, which recovers a family of stochastic policies from offline data. When coupled with zero-order search over compact policy embeddings, this algorithm can sidestep iterative optimization schemes, and optimizes general utilities directly at test-time. Across both didactic and high-dimensional experiments, we demonstrate that our method retains favorable properties of FB algorithms, while also extending their range to more general RL problems.