Unified Policy Value Decomposition for Rapid Adaptation

Cristiano Capone; Luca Falorsi; Andrea Ciardiello; Luca Manneschi

cs.LGMar 18, 2026

Unified Policy Value Decomposition for Rapid Adaptation

Authors: Cristiano Capone, Luca Falorsi, Andrea Ciardiello, Luca Manneschi

Organizations: Computational Neuroscience Unit, Istituto Superiore di Sanit`a, 00161, Rome, Italy · Ospedale Santa Lucia, Rome, Italy · School of Computer Science, University of Sheffield,Sheffield,S10 2TN, United Kingdom

Abstract

Rapid adaptation in complex control systems remains a central challenge in reinforcement learning. We introduce a framework in which policy and value functions share a low-dimensional coefficient vector - a goal embedding - that captures task identity and enables immediate adaptation to novel tasks without retraining representations. During pretraining, we jointly learn structured value bases and compatible policy bases through a bilinear actor-critic decomposition. The critic factorizes as Q = sum_k G_k(g) y_k(s,a), where G_k(g) is a goal-conditioned coefficient vector and y_k(s,a) are learned value basis functions. This multiplicative gating - where a context signal scales a set of state-dependent bases - is reminiscent of gain modulation observed in Layer 5 pyramidal neurons, where top-down inputs modulate the gain of sensory-driven responses without altering their tuning. Building on Successor Features, we extend the decomposition to the actor, which composes a set of primitive policies weighted by the same coefficients G_k(g). At test time the bases are frozen and G_k(g) is estimated zero-shot via a single forward pass, enabling immediate adaptation to novel tasks without any gradient update. We train a Soft Actor-Critic agent on the MuJoCo Ant environment under a multi-directional locomotion objective, requiring the agent to walk in eight directions specified as continuous goal vectors. The bilinear structure allows each policy head to specialize to a subset of directions, while the shared coefficient layer generalizes across them, accommodating novel directions by interpolating in goal embedding space. Our results suggest that shared low-dimensional goal embeddings offer a general mechanism for rapid, structured adaptation in high-dimensional control, and highlight a potentially biologically plausible principle for efficient transfer in complex reinforcement learning systems.

Figures & tables

Figure 1: MLP-based bilinear actor–critic architecture. A. Scheme of our architecture: the actor and critic are decomposed into $K$ parallel basis modules, policy primitives $Y_{k}(s)$ and value components $\phi_{k}(s,a)$ . B–C. Comparison between the learning curves of a traditional architecture (2-layer MLP) and a single-layer MLP with bilinear decomposition. (inset) Scheme of the navigation task: a robot with 8 DOF is asked to move in a specific direction. D. Comparison of learning curves between the cases in which the latent space $G_{k}$ is independent or shared between actor and critic. E–F. Direction encoding for actor and critic. G–I. Reward, correlation between actor and critic $G$ , and direction encoding in the $G$ space, as functions of training steps.

Figure 2: Zero-shot learning. A. Task scheme: the MuJoCo Ant agent is pre-trained on target directions and tested on new ones. Pretrained bilinear agent is evaluated on unseen goal directions (or task descriptors) without any parameter update, by conditioning on $g$ . B. Performance compared against baselines when switching to novel directions. C–D. Behavior trajectories for training and test directions, respectively, illustrating successful generalization to intermediate angles not explicitly trained. E. Direct comparison between train and test directions (averaged over trials).

Figure 3: Bilinear decomposition allows for interpretability and generalization. A. Manipulating individual gating coordinates $G_{k}$ produces consistent, semantically meaningful changes in behavior, indicating interpretable control axes. (Top) Average movement direction as a function of latent-space direction for three different amplitudes; (bottom) speed distribution for the same three amplitudes (color-coded). Notably, speed modulation emerges spontaneously, although training objectives were defined only over movement direction. B. Online adaptation of $G_{k}$ , allowing real-time solution of the current task. Reward as a function of the target direction (blue). As a reference, a task with negative reward is shown (orange).

Zero-Shot RL

See all 17 papers

Reward as Observation: Learning Reward-Based Policies for Rapid Adaptation
Sep 30, 2026Morgan Byrd, Jacob Blevins, Maks Sorokin +2Robot Policy Learning Robot Policy Adaptation
Georgia Institute of Technology, Atlanta, GA, 30308, USA · Georgia Tech Research Institute, Atlanta, GA, 30308, USA
Structured Representation Learning with Locally Linear Embeddings and Adaptive Feature Fusion
Jun 16, 2026Somjit Nath, Jackson J Cone, Derek Nowrouzezahrai +1Disentangled Representation Learning State Representation Learning
Department of Electrical and Computer Engineering McGill University Mila – Quebec AI Institute · Department of Psychology Department of Biomedical Engineering University of Calgary Canada Research Chair in Computational Behavioural Neuroscience (Tier II) · Department of Electrical and Software Engineering Schulich School of Engineering University of Calgary Mila – Quebec AI Institute Canada CIFAR AI Chair
An Agency-Transferring Model-Free Policy Enhancement Technique
Jun 8, 2026Anton Bolychev, Georgiy Malaniya, Sinan Ibrahim +1RL Control Transfer Learning
Center for Engineering Systems and Sciences · Central University; Center for Engineering Systems and Sciences; Sirius University of Science and Technology

Sep 30, 2026J/K move · Enter open · S save

Reward as Observation: Learning Reward-Based Policies for Rapid Adaptation

Morgan Byrd, Jacob Blevins, Maks Sorokin, Robert Wright+1

Georgia Institute of Technology, Atlanta, GA, 30308, USA · Georgia Tech Research Institute, Atlanta, GA, 30308, USA

Robot Policy Learning Robot Policy Adaptation Transfer Learning Zero-Shot Policy Transfer Policy Learning Generalization in RL+1

Open paper AskPDF arXiv

This paper explores a reward-based policy to achieve zero-shot transfer between source and target environments with completely different observation spaces. While humans can demonstrate impressive adaptation capabilities, deep neural network policies often struggle to adapt to a new environment and require a considerable amount of samples for successful transfer. Instead, we propose a novel reward-based policy only conditioned on rewards and actions, enabling zero-shot adaptation to new environments with completely different observations. We discuss the challenges and feasibility of a reward-based policy and then propose a practical algorithm for training. We demonstrate that a reward policy can be trained within three different environments, Pointmass, Cartpole, and 2D Car Racing, and transferred to completely different observations, such as different color palettes or 3D rendering, or Stretch robot navigation in Habitat-Sim, in a zero-shot manner. We also demonstrate that a reward-based policy can further guide the training of an observation-based policy in the target environment.

Abstract

Zero-Shot RL

Explore similar work

Reward as Observation: Learning Reward-Based Policies for Rapid Adaptation

Structured Representation Learning with Locally Linear Embeddings and Adaptive Feature Fusion

An Agency-Transferring Model-Free Policy Enhancement Technique