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
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 Yk(s) and value components ϕ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 Gk 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 Gk 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 Gk , 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).
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.
Morgan Byrd, Jacob Blevins, Maks Sorokin +2
Georgia Institute of Technology, Atlanta, GA, 30308, USA · Georgia Tech Research Institute, Atlanta, GA, 30308, USA
Neuroscientific research has revealed that the brain encodes complex behaviors by leveraging structured, low-dimensional manifolds and dynamically fusing multiple sources of information through adaptive gating mechanisms. Inspired by these principles, we propose a novel reinforcement learning (RL) framework that encourages the disentanglement of dynamics-specific and reward-specific features, drawing direct parallels to how neural circuits separate and integrate information for efficient decision-making. Our approach leverages locally linear embeddings (LLEs) to capture the intrinsic, locally linear structure inherent in many environments, mirroring the local smoothness observed in neural population activity, while concurrently deriving reward-specific features through the standard RL objective. An attention mechanism, analogous to cortical gating, adaptively fuses these complementary representations on a per-state basis. Experimental results on benchmark tasks demonstrate that our method, grounded in neuroscientific principles, improves learning efficiency and overall performance compared to conventional RL approaches, highlighting the benefits of explicitly modeling local state structures and adaptive feature selection as observed in biological systems.
Somjit Nath, Jackson J Cone, Derek Nowrouzezahrai +1
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
Training reinforcement learning (RL) policies from scratch is costly: it requires careful reward and environment design, extensive tuning, and substantial computation. Yet many control problems already have a functional but suboptimal policy available as a baseline. This paper proposes a method for embedding such a baseline into the RL training process, simultaneously improving training efficiency relative to from-scratch methods and producing a learning policy that outperforms the baseline. At each step, the method arbitrates between the baseline policy and a trainable learning policy, initially relying strongly on the baseline policy and then progressively transferring agency to the learning policy. By the end of training, the learning policy is a standalone neural network that operates without baseline policy support. The paper formalizes what it means for the baseline policy to be functional: under this policy, the agent reaches a goal set and remains there with high probability. The proposed arbitration mechanism is designed to exploit this property during training, yielding high goal-reaching rates right from the beginning of training. A theoretical analysis provides a formal interpretation of this behavior under stated assumptions and extends it to the final baseline-free regime, where explicit lower bounds are derived for the goal-reaching probability of the standalone learning policy. Empirical results on continuous-control benchmarks show that the proposed method achieves returns that match or exceed those of competitive approaches, while maintaining the highest goal-reaching rates throughout training among the compared methods -- including in the final stage, where the learning policy operates without any baseline support.
Anton Bolychev, Georgiy Malaniya, Sinan Ibrahim +1
Center for Engineering Systems and Sciences · Central University; Center for Engineering Systems and Sciences; Sirius University of Science and Technology