Soft Forward-Backward Representations for Zero-shot Reinforcement Learning with General Utilities

2026-02-06

AI Summary

This work addresses the challenge of optimizing general utility functions—such as those based on distribution matching or pure exploration, which are non-additive and not expressible as standard rewards—in zero-shot reinforcement learning. The authors propose the Maximum Entropy Soft Forward-Backward algorithm, which extends the forward-backward framework to general utility reinforcement learning for the first time. Their method learns a family of policies from offline data and, at test time, directly maximizes any differentiable utility function over state-action occupancy measures via zeroth-order optimization, without requiring task-specific iterative training. Experiments demonstrate that the approach maintains strong zero-shot generalization while effectively handling non-additive objectives in high-dimensional and pedagogical settings, significantly enhancing both flexibility and performance in general utility optimization.

Technology Category

• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Stochastic Optimization
• Machine Learning: Online Learning & Bandits

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• zero-shot reinforcement learning
• general utilities
• occupancy measure
• forward-backward algorithms
• offline RL

Innovation

• zero-shot reinforcement learning
• general utilities
• forward-backward algorithm
• maximum entropy
• offline RL