Do as the Romans Do: Learning Universal Behaviors from Heterogeneous Agents

# Do as the Romans Do: Learning Universal Behaviors from Heterogeneous Agents
Source: [https://arxiv.org/html/2606.18537](https://arxiv.org/html/2606.18537)
Caleb Chang1, Davin Win Kyi1, Natasha Jaques∗1, Karen Leung∗1,2 1University of Washington,2NVIDIA ∗indicates equal contribution\. Correspondence:cchang26@uw\.edu

###### Abstract

Humans often acquire new skills by observing others, since observed behaviors implicitly reveal how to act in an environment\. However, observations drawn from a heterogeneous population introduce conflicting behavioral signals, making it difficult to determine which behaviors are worth imitating\. We address this challenge with General Reward Inference and Disentanglement \(GRID\), a social learning method that extracts universally useful behaviors from a heterogeneous population of demonstrators pursuing different goals\. GRID decomposes per\-agent reward functions into a general reward, capturing behaviors shared across all agents, and specific rewards, capturing individual preferences and objectives\. Training exclusively on the general reward provides a new paradigm of generalist pretraining\. It yields a generalist agent that internalizes universal environmental competencies, such as safety and basic task proficiency, without the mode\-averaging bias that afflicts standard learning from demonstration techniques\. This generalist serves as a superior prior for fine\-tuning to downstream tasks, including preferences unseen during training\. Experiments across a synthetic basis function decomposition, multi\-agent Craftax, and a continuous autonomous driving simulator \(Highway\-Env\) confirm that GRID successfully disentangles reward structure in a semantically meaningful way, outperforms standard learning from demonstration baselines, and enables more efficient and stable specialization\.

## 1Introduction

“When in Rome, do as the Romans do”—butwhichRoman should we follow? Observed behaviors of individuals implicitly reveal information about how one should act in the environment, yet not all agents follow the same objective, and blindly imitating any one of them can be dangerous or counterproductive\. Consider a robot placed among a crowd at a traffic intersection: different pedestrians have different destinations, so they cross the street in different directions—but everyone respects the crosswalk signal\. The robot must infer that the crosswalk status is the relevant shared cue, while individual destinations and the behaviors that pursue them are not\. Or consider a robot navigating a building where a patch of floor along the shortest path is slippery\. Humans will naturally avoid this patch, taking a longer route regardless of their destination\. An intelligent agent need not know any individual’s goal to extract something valuable: the avoidance behavior is universal across all agents\. By identifying what every agent consistently does, independent of their individual objectives, a learning agent can infer what is broadly important and adopt those behaviors itself\.

The challenge is extracting these shared signals from a heterogeneous population\. Standard learning from demonstration \(LfD\) is vulnerable to potentially conflicting, multi\-modal demonstrations and demands additional mechanisms to address variations\[[31](https://arxiv.org/html/2606.18537#bib.bib896)\]\. Inverse reinforcement learning \(IRL\), a subset of LfD, recovers an unknown reward function\. When a population pursues heterogeneous objectives, standard IRL algorithms tend to average conflicting goals together, producing a suboptimal policy that fails to achieve any specific task effectively\[[20](https://arxiv.org/html/2606.18537#bib.bib902)\]\. Reinforcement learning \(RL\) sidesteps this by learning directly from environmental feedback through a known reward function\[[32](https://arxiv.org/html/2606.18537#bib.bib919)\]\. However, RL is typically sample\-inefficient and brittle to changes in the environment, leading to poor generalization capabilities\. In short, neither RL nor standard LfD gracefully handles learning from a diverse population of agents pursuing different, unknown goals, in spite of the fact that these agents may reveal extremely valuable information\.

This raises a fundamental question:*What should an agent learn from such a population*? We posit that the most valuable signals are those*shared across agents*—behaviors that are universal precisely because they reflect general competence, safety, and environmental norms, rather than any individual’s specific objective\. All pedestrians obey the crosswalk signal regardless of destination\. All skilled drivers maintain safe following distances regardless of where they are headed\. All navigators avoid the slippery floor regardless of their target\. These shared behaviors are the “Roman customs” worth adopting—general norms of the environment that any capable agent implicitly follows\.

Social learning, a combination of RL and IRL, takes a step in this direction\. In this framework, an agent observes other agents in the environment to infer relevant environment information and objectives, and uses that knowledge to inform its own behavior through its own interactive experiences with the environment\. Prior works\[[8](https://arxiv.org/html/2606.18537#bib.bib897),[2](https://arxiv.org/html/2606.18537#bib.bib895)\]explicitly disentangle shared features of the environment from agent\-specific goals, but do not isolate a shared reward signal that represents universally desirable behaviors\. To our knowledge, no existing work explicitly disentangles the underlying reward of heterogeneous agents into universally shared competencies and individual preferences\. We address this gap withGeneral Reward Inference and Disentanglement \(GRID\)\.

![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/GRID_hero.png)Figure 1:General Reward Inference and Disentanglement \(GRID\) is a social learning method that extracts universally useful behaviors from a heterogeneous population of demonstrators pursuing different goals\. The learned general behavior can be used to train a generalist agent, which can then be fine\-tuned for a new task\.GRID features a novel deep learning architecture that uses an information\-theoretic objective to simultaneously capture individual\-specific behaviors and the general behaviors common across agents\. Concretely, we assume individual reward functions, available via RL environments or IRL, express individual behaviors\. Then, we decompose each into a sum of a general reward and an individual\-specific reward\. The general reward captures a universally shared structure, what every competent agent does regardless of their personal goal\. Performing RL with this general reward yields a generalist agent that has internalized the environmental norms of the observed population\. Crucially, this agent serves as a strong prior for downstream adaptation as it already knows what to do and what to avoid and can be fine\-tuned for any specific task more efficiently than training from scratch\. Unlike traditional specialist agents that must learn fundamental physics and safety constraints from scratch for each new objective, our generalist agent already knows how to operate safely in the environment and can be efficiently fine\-tuned to specific tasks\.

Statement of Contributions\.\(1\)We introduce the GRID architecture and its information\-theoretic training objective, which jointly disentangles general and specific reward structures\.\(2\)We provide intuition on general vs specific reward structure through interpretable visualizations of the learned reward decomposition across synthetic and multi\-agent domains, confirming semantic alignment with ground\-truth reward components\.\(3\)We demonstrate that the GRID generalist outperforms mode\-averaging LfD and RL baselines in broadly applicable competencies and downstream fine\-tuning to unseen tasks\.\(4\)We validate architectural design choices through ablation studies\.

## 2Related Works

GRID is at the intersection of three research objectives: learning from heterogeneous demonstrations, reward decomposition, and generalist pretraining\.

Learning from Heterogeneous Demonstrations\.Standard LfD methods\[[33](https://arxiv.org/html/2606.18537#bib.bib914),[26](https://arxiv.org/html/2606.18537#bib.bib263),[1](https://arxiv.org/html/2606.18537#bib.bib28),[37](https://arxiv.org/html/2606.18537#bib.bib29),[23](https://arxiv.org/html/2606.18537#bib.bib513)\]and powerful generative approaches\[[11](https://arxiv.org/html/2606.18537#bib.bib531),[10](https://arxiv.org/html/2606.18537#bib.bib913)\]share a common assumption that all demonstrations originate from a single, consistent demonstrator\. However, in large datasets and in the real world, demonstrations are often heterogeneous, reflecting different skills, preferences, and goals across demonstrators\. Subsequent works address heterogeneity through mutual information maximization\[[16](https://arxiv.org/html/2606.18537#bib.bib900),[36](https://arxiv.org/html/2606.18537#bib.bib373),[29](https://arxiv.org/html/2606.18537#bib.bib915),[20](https://arxiv.org/html/2606.18537#bib.bib902)\], while others use CVAEs and variational inference for multi\-modal human behavior modeling in autonomous driving applications\[[27](https://arxiv.org/html/2606.18537#bib.bib934),[12](https://arxiv.org/html/2606.18537#bib.bib937),[25](https://arxiv.org/html/2606.18537#bib.bib975)\]\. However, all of these approaches treat heterogeneity as a modeling target—the goal is to explain variations between demonstrators\. In contrast, our approach explicitly explains the shared commonalities between demonstrators\.

Reward Decomposition\.GRID contributes to a growing interest in structured reward decomposition\.van Seijenet al\.\[[34](https://arxiv.org/html/2606.18537#bib.bib917)\]manually decomposes the total reward into independent task components and trains a dedicated agent per component, but the decomposition is entirely hand\-engineered and does not generalize across agents\.Arjona\-Medinaet al\.\[[3](https://arxiv.org/html/2606.18537#bib.bib918)\]addresses delayed rewards in temporal credit assignment by redistributing rewards along relevant state\-action pairs in a trajectory\.Juozapaitiset al\.\[[14](https://arxiv.org/html/2606.18537#bib.bib920)\]decomposes rewards into interpretable sub\-components for transparency, but does so manually and without modeling agent heterogeneity\. GRID unifies these threads: the decomposition is learned without manual reward engineering, is semantic \(general vs\. specific\), and directly supports both interpretability and downstream specialization\.

Generalist Pretraining and Specialization\.The paradigm of pretraining a generalist model and adapting it to specific downstream tasks has become central to modern machine learning and robotics\[[19](https://arxiv.org/html/2606.18537#bib.bib933)\]\.Jaqueset al\.\[[13](https://arxiv.org/html/2606.18537#bib.bib922)\]introduces KL\-control fine\-tuning to prevent catastrophic forgetting of a pretrained prior, a technique later adopted for aligning LLMs with human preferences\[[18](https://arxiv.org/html/2606.18537#bib.bib923)\]\.Finnet al\.\[[9](https://arxiv.org/html/2606.18537#bib.bib387)\]learns a parameter initialization for rapid adaptation, andRakellyet al\.\[[22](https://arxiv.org/html/2606.18537#bib.bib924)\]extends this by inferring latent task representations without task labels at test time\. In contrast, GRID instantiates this paradigm through reward structure: the generalist is trained on the shared, general reward, and specialization consists of fine\-tuning on the specific reward, which encodes individual preferences\. This grounds the prior in a semantically meaningful decomposition, making the source of generalization transparent and the adaptation target well\-defined\.

Social Learning and Successor Features\.Social learning, the acquisition of behavior by observing others in a shared environment, is enabled in part by the recognition that agents in the same environment share underlying dynamics and norms\.Ndousseet al\.\[[17](https://arxiv.org/html/2606.18537#bib.bib931)\]shows that social learning can lead to improved generalization to new environments, and recent works have attempted to use Successor Features \(SFs\)\[[4](https://arxiv.org/html/2606.18537#bib.bib899)\]to enhance social learning\[[8](https://arxiv.org/html/2606.18537#bib.bib897),[2](https://arxiv.org/html/2606.18537#bib.bib895)\]based on works proposing SFs as a way to decompose dynamics and rewards\[[6](https://arxiv.org/html/2606.18537#bib.bib921),[24](https://arxiv.org/html/2606.18537#bib.bib925)\]\. While successor feature methods are elegant, in practice, they can prove difficult to optimize without knowing the shared feature representationa priori, and perhaps for this reason have not been widely adopted\. Further, they make restrictive assumptions about the format of the shared representation\.

## 3Preliminaries

Markov Decision Process\.The environment is modeled as a Markov Decision Process \(MDP\), denoted by⟨𝒮,𝒜,R,T,γ⟩\\langle\\mathcal\{S\},\\mathcal\{A\},R,T,\\gamma\\rangle\.𝒮\\mathcal\{S\}represents the state space and𝒜\\mathcal\{A\}represents the action space\.T:𝒮×𝒜×𝒮→\[0,1\]T:\\mathcal\{S\}\\times\\mathcal\{A\}\\times\\mathcal\{S\}\\rightarrow\[0,1\]is the transition function that yields the probability of entering the next states′s^\{\\prime\}from the current statessby applying actionaa\.R​\(s,a\)R\(s,a\)is the reward signal for taking actionaain statess, and the agent’s total discounted reward over trajectoryτ=⟨s0,a0,…,sT,aT⟩\\tau=\\langle s\_\{0\},a\_\{0\},\.\.\.,s\_\{T\},a\_\{T\}\\rangleisR​\(τ\)=∑t=0Tγt​R​\(st,at\)R\(\\tau\)=\\sum\_\{t=0\}^\{T\}\\gamma^\{t\}R\(s\_\{t\},a\_\{t\}\)whereγ∈\[0,1\)\\gamma\\in\[0,1\)is the discount factor\.

Reinforcement Learning and Inverse Reinforcement Learning\.The goal of RL is to learn an optimal policyπ∗:𝒮→𝒜\\pi^\{\*\}:\\mathcal\{S\}\\rightarrow\\mathcal\{A\}that maximizes the expected future reward:π∗=arg⁡maxπ⁡𝔼τ∼π​R​\(τ\)\\pi^\{\*\}=\\arg\\max\_\{\\pi\}\\mathbb\{E\}\_\{\\tau\\sim\\pi\}R\(\\tau\)\. RL assumes access to the reward functionRR\. In contrast, IRL operates within a MDP\\R\\backslash Rand receives a dataset ofNNdemonstrations𝒟=\{τ1,…​τN\}\\mathcal\{D\}=\\\{\\tau\_\{1\},\.\.\.\\tau\_\{N\}\\\}\. Its goal is to recover the underlying reward functionRRthat explains the observed behavior\. In our problem setting, we observe a population ofPPdemonstrators acting in an environment\. Each demonstrator,pp, acts according to the reward functionR\(p\)R^\{\(p\)\}that reflects their individual goals and preferences and shared common rewards\. We assume we have access toR\(p\)R^\{\(p\)\}through the RL environment or through IRL\.

## 4General Reward Inference and Disentanglement

We present the GRID framework \(Fig\.[1](https://arxiv.org/html/2606.18537#S1.F1)\) and its architecture \(Fig\.[2](https://arxiv.org/html/2606.18537#S4.F2)\)\. GRID trains a generalist agent by extracting reward signals that are shared across a heterogeneous population of demonstrators acting in an environment\. First, we build the dataset of tuples \(agent ID, state, action, reward\) or\(p,s,a,R\(p\)\)\(p,s,a,R^\{\(p\)\}\)from surrounding demonstrators\. Second, GRID disentanglesR\(p\)R^\{\(p\)\}into a general rewardRg​\(s,a\)R\_\{g\}\(s,a\)and specific rewardRs​\(s,a,ω\(p\)\)R\_\{s\}\(s,a,\\omega^\{\(p\)\}\), capturing preferences of individual agents\. Optimizing a policy onRgR\_\{g\}yields a generalist agent\. Third, the generalist agent adapts to new tasks by fine\-tuning on task\-specific rewards\.

### 4\.1Dataset Requirements

Effective disentanglement relies on a dataset of\(p,s,a,R\(p\)\)\(p,s,a,R^\{\(p\)\}\)tuples that covers both the shared environmental contexts giving rise to general behaviors and the individual\-specific contexts that distinguish agents from one another\. Without sufficient state\-action coverage across demonstrators, GRID may not reliably separate general and specific reward components\.

One practical approach to satisfying this requirement is to collect data from agents trained via multi\-agent reinforcement learning \(MARL\)\. Early in training, agents explore broadly, generating diverse state\-action coverage\. As training progresses and policies improve, agents visit more task\-relevant regions of the state space, enriching the dataset with higher\-quality behavioral examples\. Sampling trajectories and their associated reward signals throughout training, rather than only at convergence, ensures that both exploratory and near\-optimal behaviors are represented, providing the behavioral diversity GRID requires\.

### 4\.2GRID Architecture

![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/GRID_architecture.png)Figure 2:The model architecture of GRID\. The summed outputs of the general and specific Reward models yield the total reward\. The embedding network recreates the original personalized embedding input to maximize mutual information\. Losses are jointly optimized\.The GRID architecture \(Fig\.[2](https://arxiv.org/html/2606.18537#S4.F2)\) features two neural network blocks,θg\\theta\_\{g\}andθs\\theta\_\{s\}, that representRg​\(s,a\)R\_\{g\}\(s,a\)andRs​\(s,a,ω\(p\)\)R\_\{s\}\(s,a,\\omega^\{\(p\)\}\), whose outputs are summed to reconstruct the total reward:R^\(p\)​\(s,a,ω\(p\)\)=Rg​\(s,a\)\+Rs​\(s,a,ω\(p\)\)\\hat\{R\}^\{\(p\)\}\(s,a,\\omega^\{\(p\)\}\)=R\_\{g\}\(s,a\)\+R\_\{s\}\(s,a,\\omega^\{\(p\)\}\)\. The architecture is jointly trained to satisfy three requirements:\(R1\)The specific rewardRs​\(s,a,ω\(p\)\)R\_\{s\}\(s,a,\\omega^\{\(p\)\}\)captures behaviors relevant to corresponding agentpp\.\(R2\)The general reward captures behaviors that are meaningful to*all*agents\.\(R3\)It must accurately reproduce the original rewardR\(p\)R^\{\(p\)\}\. That is,R\(p\)≈R^\(p\)R^\{\(p\)\}\\approx\\hat\{R\}^\{\(p\)\}\. We address each of these requirements by jointly optimizing three loss functions:*embedding*,*emphasis*, and*reconstruction*, as described below\.

Embedding Loss\.To meet Requirement R1, we introduce a*learned*personal embedding vectorω\(p\)∈ℝd\\omega^\{\(p\)\}\\in\\mathbb\{R\}^\{d\}for each agent, whereddis the chosen embedding dimension\. Intuitively, if a personal embedding meaningfully encodes an agent’s behavior, the specific reward should be able to reconstruct it\. This ensures that similar embeddings correspond to similar reward structures\. The embedding is optimized to capture individual characteristics by maximizing the mutual informationIIbetweenω\(p\)\\omega^\{\(p\)\}and the observed behavior\(s,a,R\(p\)\)\(s,a,R^\{\(p\)\}\)\. That is, we seek to learnω\(p\)\\omega^\{\(p\)\}bymaxω\(p\)⁡I​\(ω\(p\);s,a,R\(p\)\),where​I​\(ω\(p\);s,a,R\(p\)\)=H​\(ω\(p\)\)−H​\(ω\(p\)\|s,a,R\(p\)\)\\max\_\{\\omega^\{\(p\)\}\}I\(\\omega^\{\(p\)\};s,a,R^\{\(p\)\}\),\\\>\\text\{where\}\\\>I\(\\omega^\{\(p\)\};s,a,R^\{\(p\)\}\)=H\(\\omega^\{\(p\)\}\)\-H\(\\omega^\{\(p\)\}\|s,a,R^\{\(p\)\}\)\. Unfortunately, calculatingH​\(ω\(p\)\|s,a,R\(p\)\)H\(\\omega^\{\(p\)\}\|s,a,R^\{\(p\)\}\)is intractable because it requires access to the posteriorP​\(ω\(p\)\|s,a,R\(p\)\)P\(\\omega^\{\(p\)\}\|s,a,R^\{\(p\)\}\), which assumes that we already know how to quantify the behaviors of surrounding agents\. Instead, we maximize the Evidence Lower Bound \(Eq\. \([1](https://arxiv.org/html/2606.18537#S4.E1)\)\) using an auxillary embedding networkq​\(ω\(p\)∣s,a,R\(p\)\)q\(\\omega^\{\(p\)\}\\mid s,a,R^\{\(p\)\}\)that approximatesP​\(ω\(p\)\|s,a,R\(p\)\)P\(\\omega^\{\(p\)\}\|s,a,R^\{\(p\)\}\)\[[5](https://arxiv.org/html/2606.18537#bib.bib903),[16](https://arxiv.org/html/2606.18537#bib.bib900)\]\.

I​\(ω\(p\);s,a,R\(p\)\)=H​\(ω\(p\)\)−H​\(ω\(p\)\|s,a,R\(p\)\)≥𝔼​\[log⁡\(q​\(ω\(p\)\|s,a,R\(p\)\)\)\]\+H​\(ω\(p\)\)I\(\\omega^\{\(p\)\};s,a,R^\{\(p\)\}\)=H\(\\omega^\{\(p\)\}\)\-H\(\\omega^\{\(p\)\}\|s,a,R^\{\(p\)\}\)\\geq\\mathbb\{E\}\[\\log\(q\(\\omega^\{\(p\)\}\|s,a,R^\{\(p\)\}\)\)\]\+H\(\\omega^\{\(p\)\}\)\(1\)In practice,ω\(p\)∼𝒩​\(0,Id\)\\omega^\{\(p\)\}\\sim\\mathcal\{N\}\(0,I\_\{d\}\)is sampled once before training\. Then, during training, the embedding network,qq, reconstructs the originalω\(p\)\\omega^\{\(p\)\}input from\(s,a,R^\(p\)\)\(s,a,\\hat\{R\}^\{\(p\)\}\)\. Maximizing the ELBo is equivalent to minimizing Eq\. \([2](https://arxiv.org/html/2606.18537#S4.E2)\) and backpropagating the gradients into the originalω\(p\)\\omega^\{\(p\)\}\[[28](https://arxiv.org/html/2606.18537#bib.bib901),[29](https://arxiv.org/html/2606.18537#bib.bib915)\]\. During the course of training, the inputω\(p\)\\omega^\{\(p\)\}and the outputω^\(p\)\\hat\{\\omega\}^\{\(p\)\}converge to each other’s value\.

ℒembedding=1d​∑i=1d\(ω^i\(p\)−ωi\(p\)\)2,where​q​\(s,a,R^\(p\)\)=ω^\(p\)\\mathcal\{L\}\_\{\\text\{embedding\}\}=\\frac\{1\}\{d\}\\sum\_\{i=1\}^\{d\}\(\\hat\{\\omega\}\_\{i\}^\{\(p\)\}\-\\omega\_\{i\}^\{\(p\)\}\)^\{2\},\\\>\\text\{ where \}\\\>q\(s,a,\\hat\{R\}^\{\(p\)\}\)=\\hat\{\\omega\}^\{\(p\)\}\(2\)Emphasis Loss\.To meet Requirement R2, we prevent the trivial solution in whichRg=0R\_\{g\}=0for all states and actions\. Since the per agent total reward,R\(p\)​\(s,a,ω\(p\)\)R^\{\(p\)\}\(s,a,\\omega^\{\(p\)\}\), and the corresponding specific reward,Rs​\(s,a,ω\(p\)\)R\_\{s\}\(s,a,\\omega^\{\(p\)\}\), share identical neural network inputs,RsR\_\{s\}can representR\(p\)R^\{\(p\)\}in totality, that is, whenRg=0R\_\{g\}=0,Rs​\(s,a,ω\(p\)\)=R\(p\)​\(s,a,ω\(p\)\)R\_\{s\}\(s,a,\\omega^\{\(p\)\}\)=R^\{\(p\)\}\(s,a,\\omega^\{\(p\)\}\)\. We address this with an emphasis loss that incentivizesRgR\_\{g\}to account for the majority of the total reward, leavingRsR\_\{s\}to explain only the residual individual\-specific component\. To further reinforce this asymmetry,RsR\_\{s\}is intentionally parameterized with fewer weights thanRgR\_\{g\}\. As such, we minimize the following ratio:

ℒemphasis=−exp⁡\(Rg​\(s,a\)\)exp⁡\(Rg​\(s,a\)\)\+exp⁡\(Rs​\(s,a,ω\(p\)\)\)\\mathcal\{L\}\_\{\\text\{emphasis\}\}=\\frac\{\-\\exp\(R\_\{g\}\(s,a\)\)\}\{\\exp\(R\_\{g\}\(s,a\)\)\+\\exp\(R\_\{s\}\(s,a,\\omega^\{\(p\)\}\)\)\}\(3\)
Reconstruction Loss\.To meet Requirement R3, we minimize the following squared error loss between the predicted total rewardR^\(p\)\\hat\{R\}^\{\(p\)\}and the original rewardR\(p\)R^\{\(p\)\}for faithful decomposition of the general and specific components to the observed reward signal\.

ℒreconstruction=\(R^\(p\)​\(s,a,ω\(p\)\)−R\(p\)\)2\\mathcal\{L\}\_\{\\text\{reconstruction\}\}=\(\\hat\{R\}^\{\(p\)\}\(s,a,\\omega^\{\(p\)\}\)\-R^\{\(p\)\}\)^\{2\}\(4\)
Altogether, we jointly optimize these loss functions with gradient descent to disentangle general and specific rewards:ℒtotal=ℒreconstruction\+λ1​ℒemphasis\+λ2​ℒembedding\\mathcal\{L\}\_\{\\text\{total\}\}=\\mathcal\{L\}\_\{\\text\{reconstruction\}\}\+\\lambda\_\{1\}\\mathcal\{L\}\_\{\\text\{emphasis\}\}\+\\lambda\_\{2\}\\mathcal\{L\}\_\{\\text\{embedding\}\}whereλ1,λ2\>0\\lambda\_\{1\},\\lambda\_\{2\}\>0\.

### 4\.3Generalist and Specialist Training

Generalist Training\.We train a generalist agent by applying standard RL to the general reward componentRg​\(st,at\)R\_\{g\}\(s\_\{t\},a\_\{t\}\)learned by GRID\. A key practical consideration arises from the episodic structure of RL training\. GRID assumes reward decomposition is additive, so for any constantcc,Rg\+cR\_\{g\}\+candRs−cR\_\{s\}\-csatisfy\(Rg\+c\)\+\(Rs−c\)=Rg\+Rs\(R\_\{g\}\+c\)\+\(R\_\{s\}\-c\)=R\_\{g\}\+R\_\{s\}, preserving the total reward while shifting mass between the two components\. A constant offset in the general reward model’s output is not penalized and can arise during training\. However, this offset has consequences in episodic settings\. When a constantccis present in the reward function, the total discounted return over an episode with lengthTTacquires an additional term that scales with episode duration:R​\(τ\)=∑t=0Tγt​\(R​\(st,at\)\+c\)=∑t=0Tγt​R​\(st,at\)\+c​1−γT1−γR\(\\tau\)=\\sum^\{T\}\_\{t=0\}\\gamma^\{t\}\(R\(s\_\{t\},a\_\{t\}\)\+c\)=\\sum^\{T\}\_\{t=0\}\\gamma^\{t\}R\(s\_\{t\},a\_\{t\}\)\+c\\frac\{1\-\\gamma^\{T\}\}\{1\-\\gamma\}\. This extra term acts as an implicit survival incentive whose sign is determined bycc\. Ifc\>0c\>0, the agent is rewarded for prolonging the episode; ifc<0c<0, the agent incurs a cumulative penalty for each additional timestep, creating an incentive to terminate the episode prematurely\. In any environment where the agent has agency over episode termination, eg\. through failure states, goal\-reaching, or other terminal transitions, a negative offset can cause the agent to seek early termination rather than learning the intended behavior\. To eliminate this side effect, we rescale the output of the general reward model to\[−1,1\]\[\-1,1\]before training\. Since GRID captures the general reward structure, this rescaling ensures that poor behaviors are penalized and ideal behaviors are rewarded\.

Specialist Training\.Once the generalist is trained, we fine\-tune the generalist policy to create a specialist agent\. We use\[[13](https://arxiv.org/html/2606.18537#bib.bib922)\]’s KL\-control method to fine\-tune the generalist policyπg\\pi\_\{g\}for a new task defined by a new specialist reward functionRS​\(st,at\)R\_\{S\}\(s\_\{t\},a\_\{t\}\)\. The fine\-tuning objective maximizes the reward while penalizing divergence from the generalist priorπg\\pi\_\{g\}:ℒ​\(πs\)=𝔼π​\[∑tRs​\(st,at\)\+β​\(log⁡πg​\(at∣st\)−log⁡πs​\(at∣st\)\)\]\\mathcal\{L\}\(\\pi\_\{s\}\)=\\mathbb\{E\}\_\{\\pi\}\\left\[\\sum\_\{t\}R\_\{s\}\(s\_\{t\},a\_\{t\}\)\+\\beta\\big\(\\log\\pi\_\{g\}\(a\_\{t\}\\mid s\_\{t\}\)\-\\log\\pi\_\{s\}\(a\_\{t\}\\mid s\_\{t\}\)\\big\)\\right\]\. The KL penalty prevents catastrophic forgetting of the general competencies encoded inπg\\pi\_\{g\}, while allowing the specialist policy to deviate toward agent\-specific preferences\. This mirrors the same use of fine\-tuning in language models\[[18](https://arxiv.org/html/2606.18537#bib.bib923)\], but grounds the prior in a semantically decomposed reward signal rather than a pretraining corpus, making the behavioral boundary between generalist and specialist explicit and interpretable\.

## 5Experiments

We outline our experimental setup to evaluate General Reward Inference and Disentanglement \(GRID\)\. Our evaluation aims to answer two central questions:

Q1: Can GRID successfully disentangle general and specific reward structures in a semantically meaningful and interpretable way?We hypothesize that GRID can isolate a general reward structure that captures universal capabilities from specific reward structures that encode individual agents’ specific behaviors\. Since evaluating semantic meaning in sequential decision\-making can be subjective, we conduct experiments in domains where the true underlying reward components can be clearly visualized and compared against GRID’s outputs\. We further validate the architectural design choices through ablation studies on the specific model capacity asymmetry and emphasis loss\.

Q2: Does training on the extracted general reward produce a robust generalist agent that can be efficiently adapted to new tasks?We train a PPO agent\[[30](https://arxiv.org/html/2606.18537#bib.bib926)\]exclusively on the learned general reward and evaluate its generalist and fine\-tuned performance against standard LfD baselines trained on heterogeneous data\. We hypothesize that, unlike the LfD baselines, which are prone to mode\-averaging behavior under conflicting demonstrations, the GRID generalist will exhibit unbiased, generally competent behavior, and it will serve as a superior prior for downstream fine\-tuning\.

### 5\.1Experimental Setup

We evaluate GRID across three domains of increasing complexity\. For each domain we describe the environment, the reward structure, and the specific comparisons made\.

![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/basis-func.png)Figure 3:Basis Function Decomposition\.Basis Function Decomposition\.A synthetic continuous domain designed to provide clear intuition for reward disentanglement\. We define three total reward functions:R1​\(x\)=cos⁡\(x\)\+xR\_\{1\}\(x\)=\\cos\(x\)\+x,R2​\(x\)=cos⁡\(x\)R\_\{2\}\(x\)=\\cos\(x\), andR3​\(x\)=cos⁡\(x\)−xR\_\{3\}\(x\)=\\cos\(x\)\-x\. GRID is trained to predict these targets, where the shared cosine component represents the ground\-truth general reward and the linear terms\(\+x,0,−x\)\(\+x,0,\-x\)represent ground\-truth specific rewards\. No baselines are evaluated in this domain, as its purpose is purely to provide a visualizable proof of disentanglement\.

![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/edit_craftax_divided_2.png)Figure 4:Craftax Domain\. We displaymaxa⁡R​\(s,a\)\\max\_\{a\}R\(s,a\)for total, specific, and general rewards over the state space for each agent\.Craftax\.An8×88\\times 8multi\-agent discrete world, representing a simplified version of the Minecraft game\[[35](https://arxiv.org/html/2606.18537#bib.bib928)\]\. Three agents each pursue a unique objective: collect wood, stone, or diamond, and all share a penalty for stepping on lava\. This domain tests GRID’s ability to disentangle a shared safety constraint from different individual goals in a discrete state\-action space\. As with the Basis Function domain, this Craftax environment offers an interpretable visualization of learned reward maps for qualitative verification\.

![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/highway_speed_lane.png)

![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/highway-reward.png)

Figure 5:Highway\-Env Domain\. The learned specific reward increases when the agent moves to its preferred lane\. The learned general reward increases when the agent accelerates, as evidenced by the greater spacing in the trajectory trail\. Other vehicle trajectories are not visualized for visual clarity\.Highway\-Env\.A continuous state and discrete action space autonomous driving simulator, modified into a three\-agent setting with four lanes and additional moving vehicles\[[15](https://arxiv.org/html/2606.18537#bib.bib124)\]\. All agents share a general reward for maintaining high speed and avoiding crashes\. Each agent additionally holds a lane preference: Agentiiprefers Laneii, fori=0,1,2i=0,1,2\. Lane 3 is deliberately withheld from training to evaluate downstream specialization to an unseen preference\. This domain evaluates GRID’s ability to recover the high\-speed reward and collision penalty structure in the general reward\. Here, we move beyond qualitative visualizations and empirically assess the utility of developing a generalist agent from general rewards, as well as its potential for fine\-tuning on new tasks\.

Baselines and Comparison Metrics\.We compare GRID against Behavioral Cloning \(BC\)\[[33](https://arxiv.org/html/2606.18537#bib.bib914)\]and Adversarial Inverse Reinforcement Learning \(AIRL\)\[[10](https://arxiv.org/html/2606.18537#bib.bib913)\]\. We choose BC as a baseline because it is a supervised learning\-based approach that has been widely adopted in robotics \(next\-action prediction\) and language \(next\-token prediction\) domains for generalist pretraining\. We also compare GRID against AIRL because it is a state\-of\-the\-art approach in IRL due to its powerful generative model\. Following Sec\.[4\.1](https://arxiv.org/html/2606.18537#S4.SS1), we use Independent PPO\[[7](https://arxiv.org/html/2606.18537#bib.bib927)\]to create expert heterogeneous demonstrations\. We anticipate that BC and AIRL will exhibit mode\-averaging behavior over the three lane preferences, while the GRID generalist—trained only on the extracted general reward—will exhibit unbiased behavior\.

All methods are evaluated on the ground truth shared general reward using total return, episode length, and lane occupancy entropy as metrics\. An entropy of 2\.0 denotes a maximally uniform distribution across lanes; lower values indicate lane bias\. Downstream adaptability is assessed by fine\-tuning BC, AIRL, and GRID priors to prefer Lane 3 via KL\-control\. General rewards for maintaining high speed and avoiding crashes are withdrawn in the fine\-tuning environment to test each prior’s generalist capability\. These rewards are added back during the evaluation phase to assess generalist capability maintenance\. We evaluate fine\-tuning performance with a Lane 3 Absolute Occupancy metric, the number of steps an agent is in Lane 3 divided by the maximum environment steps\.

Additionally, we consider two specialist RL agents: Right Specialist and Wrong Specialist\. The Right Specialist is trained from scratch on the true general and Lane 3 preference rewards within the same step budget as fine\-tuning\. We include the Right Specialist to assess whether or not generalist pretraining is actually helpful in policy learning\. The Wrong Specialist is trained to be an expert on the wrong task \(prefers Lane 0\)\. We treat the Wrong Specialist as a prior for fine\-tuning to compare the effects of a single\-mode bias and a multi\-modal bias from heterogeneous demonstrations\. Implementation details can be found in Appendix[C](https://arxiv.org/html/2606.18537#A3)\.

![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/highway-crash-scene.png)
![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/highway-crash-reward.png)

Figure 6:General Reward for collision\.

## 6Results

Q1: Semantically Meaningful Disentanglement\.Figures[3](https://arxiv.org/html/2606.18537#S5.F3)and[4](https://arxiv.org/html/2606.18537#S5.F4)visually confirm GRID’s ability to disentangle reward components accurately\. In the Basis Function experiment \(see Fig\.[3](https://arxiv.org/html/2606.18537#S5.F3)\), the learned general reward captures the sinusoidal component \(with an offset\), while each specific reward network isolates its corresponding linear term\. The decomposition matches the ground\-truth structure, providing interpretable, verifiable evidence that GRID’s disentanglement learns meaningful disentanglements\. This result transfers to the Craftax environment\. Figure[4](https://arxiv.org/html/2606.18537#S5.F4)presents the Craftax environment \(top left\), the learned total reward for each agent \(top right\), the rewards for procuring specific materials \(bottom right\), and the shared general reward \(bottom left\)\. The general reward correctly identifies lava locations as shared negative penalties across all agents, while item\-gathering preferences \(tree wood, diamond, and stone\) are cleanly partitioned into the agent\-specific reward maps with no observable overlap\. Note that the specific reward does not necessarily coincide exactly with the trees/diamonds/stones locations, because items can be collected \(i\.e\., reward gained\) if the agent is adjacent to those items\. For the Highway\-Env experiment, we see that in Figs\.[5](https://arxiv.org/html/2606.18537#S5.F5)and[6](https://arxiv.org/html/2606.18537#S5.F6), the extracted general reward penalizes collisions and increases with higher speeds, matching the environment’s[reward structure](https://highway-env.farama.org/rewards/), thereby filtering individual lane preferences into the specific reward\. Fig\.[6](https://arxiv.org/html/2606.18537#S5.F6)demonstrates the same safety representation in the general reward, where a decrease in reward corresponds to the agent colliding with another vehicle\.

![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/best_curve_lane_frac.png)
![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/best_curve_length.png)
![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/best_curve_return.png)

Figure 7:GRID, AIRL, BC, and Wrong Specialist are fine\-tuned on Lane 3 preference only, without access to generalist rewards\. The Correct Specialist is trained from scratch with full reward access\.![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/complex-highway.png)Figure 8:Generalist agent \(green vehicle\) navigates around other cars \(blue vehicles\) at high speeds\.GeneralistSpecialistMethodReward \(↑\\uparrow\)Length \(↑\\uparrow\)Entropy \(↑\\uparrow\)Reward \(↑\\uparrow\)Length \(↑\\uparrow\)Lane 3 Abs \(↑\\uparrow\)Right Specialist–––22\.1±1\.722\.1\\pm 1\.723\.7±1\.723\.7\\pm 1\.70\.232±0\.0290\.232\\pm 0\.029Wrong Specialist–––34\.0±1\.634\.0\\pm 1\.634\.9±1\.534\.9\\pm 1\.50\.302±0\.0310\.302\\pm 0\.031BC\[[33](https://arxiv.org/html/2606.18537#bib.bib914)\]29\.8±1\.829\.8\\pm 1\.831\.7±1\.731\.7\\pm 1\.71\.931\.9329\.9±1\.829\.9\\pm 1\.832\.0±1\.732\.0\\pm 1\.70\.272±0\.0320\.272\\pm 0\.032AIRL\[[10](https://arxiv.org/html/2606.18537#bib.bib913)\]28\.8±1\.828\.8\\pm 1\.830\.4±1\.730\.4\\pm 1\.71\.891\.8924\.3±2\.024\.3\\pm 2\.026\.0±1\.926\.0\\pm 1\.90\.188±0\.0280\.188\\pm 0\.028GRID \(ours\)34\.7±0\.8\\mathbf\{34\.7\\pm 0\.8\}37\.9±0\.8\\mathbf\{37\.9\\pm 0\.8\}1\.98\\mathbf\{1\.98\}34\.9±1\.0\\mathbf\{34\.9\\pm 1\.0\}37\.9±1\.0\\mathbf\{37\.9\\pm 1\.0\}0\.349±0\.025\\mathbf\{0\.349\\pm 0\.025\}

Table 1:Generalist and Specialist Comparisons in the Highway\-Env\. See Tables[2](https://arxiv.org/html/2606.18537#A1.T2)and[3](https://arxiv.org/html/2606.18537#A1.T3)in the Appendix for the full lane distributions\. Means and standard errors are shown\.Q2: Generalization and Specialization\.GRID develops a superior generalist\. Figure[8](https://arxiv.org/html/2606.18537#S6.F8)shows its ability to navigate through traffic at high speeds\. This is validated empirically in Table[1](https://arxiv.org/html/2606.18537#S6.T1), where GRID outperforms BC and AIRL in all metrics\. The GRID generalist achieves a16\.4%−20\.5%16\.4\\%\-20\.5\\%higher return and a19\.6%−24\.7%19\.6\\%\-24\.7\\%increase in survival length compared to the baselines\. Notably, while BC and AIRL underperform compared to the original demonstrators \(surviving31\.731\.7and30\.430\.4steps compared to33\.033\.0steps of the IPPO expert\), GRID surpasses the expert average with 37\.9 steps\. This suggests that isolating the general reward removes the noise introduced by conflicting individual preferences, yielding a cleaner behavioral signal than what any single demonstrator provides\.

GRID also exhibits the most unbiased generalist behavior, achieving a near\-uniform lane distribution with entropy1\.981\.98\(max:2\.02\.0\)\. In contrast, BC and AIRL suffer from severe mode\-averaging because the training data contains agents who prefer Lanes 0, 1, and 2\. Both baselines develop an unintended bias towards Lane 1, respectively, spending on average36\.1%36\.1\\%and37\.5%37\.5\\%of the episode there\. AIRL is particularly susceptible to the multi\-modal data distribution\. Despite a BC warm\-start, the discriminator fails to fit the heterogeneous demonstrations with a single reward estimate, leading to collapse and deteriorating performance\. When fine\-tuning to prefer the previously unseen Lane 3 via KL\-control, the GRID generalist serves as a far superior prior, consistently outperforming both baselines in return, episode length, and target lane adherence \(Fig\.[7](https://arxiv.org/html/2606.18537#S6.F7), Table[1](https://arxiv.org/html/2606.18537#S6.T1)\)\. Our approach follows the new desired preference28\.3%28\.3\\%and85\.6%85\.6\\%more often than BC and AIRL, while maintaining its generalist capabilities\. In particular, GRID survives the longest with the least standard error, suggesting that GRID robustly maintains its safety without any compromise to its fine\-tuning or high\-speed performance\.

We further contextualize these results against two specialist baselines that have access to the true rewards\. GRID and BC exceed the Right Specialist in Lane 3 absolute occupancy, meeting our expectations, but AIRL does not, for the reasons discussed above\. The Wrong Specialist, trained on Lane 0 and fine\-tuned to Lane 3, reveals an insightful finding: it serves as a stronger prior than both BC and AIRL generalists\. We attribute this to the cost of resolving multi\-modal confusion\. BC and AIRL must first unlearn conflicting lane preferences accumulated from heterogeneous training data before acquiring the new preference\. In contrast, the Wrong Specialist carries only a single, coherent bias that is easier to override\. GRID outperforms the Wrong Specialist, suggesting a clear hierarchy for generalist fine\-tuning: an unbiased generalist is ideal, a single\-bias prior is preferable to a multi\-modal one, and multi\-modal bias, as exhibited by standard LfD methods, is most costly to correct\. Together, these results confirm that disentangling the general reward from individual\-specific preferences yields a generalist agent that is not only more competent and unbiased than mode\-averaging baselines but also a more effective foundation for specialization to new tasks\.

![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/ablation.png)Figure 9:Specific Model Capacity and Emphasis Loss Ablation\. Both ablations’ general rewards struggle with learning the true sinusoidal structure \(left\)\. The specific reward completely represents the total reward \(middle\)\. Some sinusoidal structure is leaked into the specific reward \(right\)\.### 6\.1Ablation Studies

We conduct an ablation study to confirm the architectural design choices that address Requirement R2 in Sec\.[2](https://arxiv.org/html/2606.18537#S4.F2)\. In the first ablation, Specific Model Capacity, we increase the specific reward model’s network size and remove the emphasis loss \(Eq\. \([3](https://arxiv.org/html/2606.18537#S4.E3)\)\) from learning\. In the second ablation, Emphasis Loss, we restore the emphasis loss, but keep the large network size\. Figure[9](https://arxiv.org/html/2606.18537#S6.F9)depicts the results: the general reward of both ablations \(left\); the specific reward for the Specific Model Capacity ablation \(middle\); and the specific reward for the Emphasis Loss ablation \(right\)\.

Specific Model Capacity\.We remove the specific reward’s capacity bottleneck by increasing its network parameter count to match the general reward’s parameter count\. The plot on the left shows a complete collapse of the general reward,Rg=0R\_\{g\}=0for allxx, and the middle plot showsRs​\(x,ω\(p\)\)=R\(p\)​\(x,ω\(p\)\)R\_\{s\}\(x,\\omega^\{\(p\)\}\)=R^\{\(p\)\}\(x,\\omega^\{\(p\)\}\)\. Without the capacity bottleneck and emphasis loss, the general reward is meaningless and no generalist agent can be trained\. More details can be found in[4](https://arxiv.org/html/2606.18537#A2.T4)\.

Emphasis Loss\.We restore the emphasis loss, but keep the high\-capacity specific reward network, resulting in some sinusoidal behavior captured by the general reward\. Its curve exhibits some local maxima and minima that correspond tocos⁡\(x\)\\cos\(x\)\. The plot on the right shows that the supposed linear specific reward actually exhibits some sinusoidal behavior as well\. Without the capacity bottleneck, the general and specific reward networks unpredictably compete to explain the total reward\. Both constraints are therefore necessary: the bottleneck enforces a hard capacity asymmetry, while the emphasis loss provides a soft inductive bias toward explanatory dominance of the general reward\.

## 7Conclusion, Limitations, and Future Work

We introduce a new paradigm of generalist pretraining via General Reward Inference and Disentanglement \(GRID\), a method for social learning that determines which behaviors are worth acquiring from a heterogeneous population of agents\. GRID successfully disentangles reward signals into a general component that captures useful universal behaviors and a specific component that encodes individual preferences\. We demonstrate that performing RL on the extracted general reward yields a generalist agent that internalizes fundamental environmental norms, outperforming LfD baselines\. However, we note that GRID relies on the availability of rewards for each demonstrator, either through the environment simulator or IRL\. When IRL is used, the quality of GRID’s disentanglement may be dependent on the quality of the recovered reward\. Furthermore, demonstrations must cover sufficient behavior coverage as described in Sec\.[4\.1](https://arxiv.org/html/2606.18537#S4.SS1)\. As such, extending GRID to a pure IRL setting is a promising direction to investigate\. Inferring both general and specific rewards directly from trajectories without access to environment reward signals would substantially broaden its applicability\. Also, scaling to larger agent populations and richer social environments, such as real\-world traffic or embodied navigation, would test the robustness of the disentanglement under greater behavioral diversity\.

## Acknowledgments

We thank Eric Ye for his assistance with the Craftax experimental environment setup and Yancheng Liang for helpful discussions\.

## References

- \[1\]\(2004\)Apprenticeship learning via inverse reinforcement learning\.InInt\. Conf\. on Machine Learning,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1)\.
- \[2\]M\. Abdulhai, N\. Jaques, and S\. Levine\(2023\)Basis for Intentions: Efficient Inverse Reinforcement Learning using Past Experience\.Available at[https://arxiv\.org/abs/2208\.04919](https://arxiv.org/abs/2208.04919)\.Cited by:[§1](https://arxiv.org/html/2606.18537#S1.p4.1),[§2](https://arxiv.org/html/2606.18537#S2.p5.1)\.
- \[3\]J\. A\. Arjona\-Medina, M\. Gillhofer, M\. Widrich, T\. Unterthiner, J\. Brandstetter, and S\. Hochreiter\(2019\)RUDDER: Return Decomposition for Delayed Rewards\.InConf\. on Neural Information Processing Systems,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p3.1)\.
- \[4\]A\. Barreto, W\. Dabney, R\. Munos, J\. J\. Hunt, T\. Schaul, H\. Hasselt, and D\. Silver\(2017\)Successor Features for Transfer in Reinforcement Learning\.InConf\. on Neural Information Processing Systems,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p5.1)\.
- \[5\]X\. Chen, Y\. Duan, R\. Houthooft, J\. Schulman, I\. Sutskever, and P\. Abbeel\(2016\)InfoGAN: Interpretable Representation Learning by Information Maximizing Generative Adversarial Nets\.InConf\. on Neural Information Processing Systems,Cited by:[§4\.2](https://arxiv.org/html/2606.18537#S4.SS2.p2.11)\.
- \[6\]P\. Dayan\(1993\)Improving Generalization for Temporal Difference Learning: The Successor Representation\.Neural Computation5,pp\. 613–624\.Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p5.1)\.
- \[7\]C\. S\. De Witt, T\. Gupta, D\. Makoviichuk, V\. Makoviychuk, P\. H\. Torr, M\. Sun, and S\. WhitesonIs independent learning all you need in the starcraft multi\-agent challenge?\.Available at[https://arxiv\.org/abs/2011\.09533](https://arxiv.org/abs/2011.09533)\.Cited by:[§5\.1](https://arxiv.org/html/2606.18537#S5.SS1.p5.1)\.
- \[8\]A\. Filos, C\. Lyle, Y\. Gal, S\. Levine, N\. Jaques, and G\. Farquhar\(2021\)PsiPhi\-Learning: Reinforcement Learning with Demonstrations using Successor Features and Inverse Temporal Difference Learning\.InInt\. Conf\. on Machine Learning,Cited by:[§1](https://arxiv.org/html/2606.18537#S1.p4.1),[§2](https://arxiv.org/html/2606.18537#S2.p5.1)\.
- \[9\]C\. Finn, P\. Abbeel, and S\. Levine\(2017\)Model\-Agnostic Meta\-Learning for Fast Adaptation of Deep Networks\.InInt\. Conf\. on Machine Learning,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p4.1)\.
- \[10\]J\. Fu, K\. Luo, and S\. Levine\(2018\)Learning Robust Rewards with Adversarial Inverse Reinforcement Learning\.InInt\. Conf\. on Learning Representations,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1),[§5\.1](https://arxiv.org/html/2606.18537#S5.SS1.p5.1),[Table 1](https://arxiv.org/html/2606.18537#S6.T1.24.24.24.7)\.
- \[11\]J\. Ho and S\. Ermon\(2016\)Generative Adversarial Imitation Learning\.InConf\. on Neural Information Processing Systems,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1)\.
- \[12\]B\. Ivanovic, K\. Leung, E\. Schmerling, and M\. Pavone\(2021\)Multimodal Deep Generative Models for Trajectory Prediction: A Conditional Variational Autoencoder Approach\.IEEE Robotics and Automation Letters6\(2\),pp\. 295–302\.Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1)\.
- \[13\]N\. Jaques, S\. Gu, D\. Bahdanau, J\. M\. Hernández\-Lobato, R\. E\. Turner, and D\. Eck\(2017\)Sequence Tutor: Conservative Fine\-Tuning of Sequence Generation Models with KL\-control\.InInt\. Conf\. on Machine Learning,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p4.1),[§4\.3](https://arxiv.org/html/2606.18537#S4.SS3.p2.5)\.
- \[14\]Z\. Juozapaitis, A\. Koul, A\. Fern, M\. Erwig, and F\. Doshi\-Velez\(2019\)Explainable reinforcement learning via reward decomposition\.InInt\. Joint Conf\. on Artificial Intelligence,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p3.1)\.
- \[15\]E\. Leurent and contributors\(2018\)An Environment for Autonomous Driving Decision\-Making\(Website\)Note:Available at[https://github\.com/eleurent/highway\-env](https://github.com/eleurent/highway-env)Cited by:[§5\.1](https://arxiv.org/html/2606.18537#S5.SS1.p4.3)\.
- \[16\]Y\. Li, J\. Song, and S\. Ermon\(2017\)InfoGAIL: Interpretable imitation learning from visual demonstrations\.InConf\. on Neural Information Processing Systems,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1),[§4\.2](https://arxiv.org/html/2606.18537#S4.SS2.p2.11)\.
- \[17\]K\. Ndousse, D\. Eck, S\. Levine, and N\. Jaques\(2020\)Emergent social learning via multi\-agent reinforcement learning\.InInternational Conference on Machine Learning,External Links:[Link](https://api.semanticscholar.org/CorpusID:235621490)Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p5.1)\.
- \[18\]L\. Ouyang, J\. Wu, X\. Jiang, D\. Almeida, C\. Wainwright, P\. L\.Mishkin, C\. Zhang, S\. Agarwal, K\. Slama, A\. Ray, J\. Schulman, J\. Hilton, F\. Kelton, L\. Miller, M\. Simens, A\. Askell, P\. Welinder, P\. Christiano, J\. Leike, and L\. R\.\(2022\)Training language models to follow instructions with human feedback\.InConf\. on Neural Information Processing Systems,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p4.1),[§4\.3](https://arxiv.org/html/2606.18537#S4.SS3.p2.5)\.
- \[19\]M\. Palan, G\. Shevchuk, N\. Landolfi, and D\. Sadigh\(2019\-06\)Learning reward functions by integrating human demonstrations and preferences\.InRobotics: Science and Systems,pp\.\.Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p4.1)\.
- \[20\]S\. Poddar, Y\. Wan, H\. Ivison, A\. Gupta, and N\. Jaques\(2024\)Personalizing Reinforcement Learning from Human Feedback with Variational Preference Learning\.InConf\. on Neural Information Processing Systems,Cited by:[§1](https://arxiv.org/html/2606.18537#S1.p2.1),[§2](https://arxiv.org/html/2606.18537#S2.p2.1)\.
- \[21\]A\. Raffin, A\. Hill, A\. Gleave, A\. Kanervisto, M\. Ernestus, and N\. Dormann\(2021\)Stable\-baselines3: reliable reinforcement learning implementations\.Journal of Machine Learning Research22\(268\),pp\. 1–8\.External Links:[Link](http://jmlr.org/papers/v22/20-1364.html)Cited by:[§B\.3](https://arxiv.org/html/2606.18537#A2.SS3.p2.1),[Appendix C](https://arxiv.org/html/2606.18537#A3.p1.1)\.
- \[22\]K\. Rakelly, A\. Zhou, D\. Quillen, C\. Finn, and S\. Levine\(2019\)Efficient Off\-Policy Meta\-Reinforcement Learning via Probabilistic Context Variables\.InInt\. Conf\. on Machine Learning,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p4.1)\.
- \[23\]D\. Ramachandran and E\. Amir\(2007\)Bayesian Inverse Reinforcement Learning\.InInt\. Joint Conf\. on Artificial Intelligence,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1)\.
- \[24\]C\. Reinke and X\. Alameda\-Pineda\(2023\)Successor Feature Representations\.Transaction on Machine Learning Research\.Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p5.1)\.
- \[25\]I\. Remy, C\. Chang, and K\. Leung\(2026\)Learning Probabilistic Responsibility Allocations for Multi\-Agent Interactions\.Available at[https://arxiv\.org/abs/2604\.13128](https://arxiv.org/abs/2604.13128)\.Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1)\.
- \[26\]S\. Ross, G\. J\. Gordon, and A\. Bagnell\(2011\)A Reduction of Imitation Learning and Structured Prediction to No\-Regret Online Learning\.InAI & Statistics,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1)\.
- \[27\]E\. Schmerling, K\. Leung, W\. Vollprecht, and M\. Pavone\(2018\)Multimodal Probabilistic Model\-Based Planning for Human\-Robot Interaction\.InProc\. IEEE Conf\. on Robotics and Automation,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1)\.
- \[28\]M\. L\. Schrum, E\. Hedlund, and M\. C\. Gombolay\(2021\)Improving Robot\-Centric Learning from Demonstration via Personalized Embeddings\.InAAAI Fall Symposium,Cited by:[§4\.2](https://arxiv.org/html/2606.18537#S4.SS2.p2.18)\.
- \[29\]M\. L\. Schrum, E\. Sumner, M\. C\. Gombolay, and A\. Best\(2024\)Maveric: A data\-driven approach to personalized autonomous driving\.IEEE Transactions on Robotics40,pp\. 1952–1965\.Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1),[§4\.2](https://arxiv.org/html/2606.18537#S4.SS2.p2.18)\.
- \[30\]J\. Schulman, F\. Wolski, P\. Dhariwal, A\. Radford, and O\. KlimovProximal policy optimization algorithms\.Available at[https://arxiv\.org/abs/1707\.06347](https://arxiv.org/abs/1707.06347)\.Cited by:[§5](https://arxiv.org/html/2606.18537#S5.p3.1)\.
- \[31\]E\. Seraj, KM\. Lee, Z\. Zaidi, Q\. Xiao, Z\. Li, A\. Nascimento, S\. Van Waveren, P\. Tambwekar, R\. Paleja, D\. Das, and M\.C\. Gombolay\(2024\)Interactive and Explainable Robot Learning: A Comprehensive Review\.Foundations and Trends in Robotics12\(2–3\),pp\. 75–349\.Cited by:[§1](https://arxiv.org/html/2606.18537#S1.p2.1)\.
- \[32\]R\. S\. Sutton and A\. G\. Barto\(2018\)Reinforcement Learning: An Introduction\.MIT Press\.Cited by:[§1](https://arxiv.org/html/2606.18537#S1.p2.1)\.
- \[33\]F\. Torabi, G\. Warnell, and P\. Stone\(2018\)Behavioral Cloning from Observation\.InInt\. Joint Conf\. on Artificial Intelligence,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1),[§5\.1](https://arxiv.org/html/2606.18537#S5.SS1.p5.1),[Table 1](https://arxiv.org/html/2606.18537#S6.T1.18.18.18.7)\.
- \[34\]H\. van Seijen, M\. Fatemi, J\. Romoff, R\. Laroche, T\. Barnes, and J\. Tsang\(2017\)Hybrid Reward Architecture for Reinforcement Learning\.InConf\. on Neural Information Processing Systems,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p3.1)\.
- \[35\]E\. Ye and N\. Jaques\(2024\)An Efficient Open World Benchmark for Multi\-Agent Reinforcement Learning\.InConf\. on Neural Information Processing Systems \- Workshop on Open\-World Agents,Cited by:[§B\.2](https://arxiv.org/html/2606.18537#A2.SS2.p2.1),[§5\.1](https://arxiv.org/html/2606.18537#S5.SS1.p3.1)\.
- \[36\]L\. Yu, T\. Yu, C\. Finn, and S\. Ermon\(2019\)Meta\-Inverse Reinforcement Learning with Probabilistic Context Variables\.InConf\. on Neural Information Processing Systems,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1)\.
- \[37\]B\. D\. Ziebart, A\. Maas, J\. A\. Bagnell, and A\. K\. Dey\(2008\)Maximum entropy inverse reinforcement learning\.InProc\. AAAI Conf\. on Artificial Intelligence,Cited by:[§2](https://arxiv.org/html/2606.18537#S2.p2.1)\.

## Appendix AAdditional Highway Results

All methods are evaluated on ground truth generalist rewards across 50 seeds\. The reward follows the[highway\-env documentation](https://highway-env.farama.org/rewards/)formula\. The Lane Occupancy metric is the ratio of steps spent in a lane to the total steps taken by the agent\.

MethodReward \(↑\\uparrow\)EpisodeLength \(↑\\uparrow\)LaneEntropy \(↑\\uparrow\)Lane OccupancyLane 0Lane 1Lane 2Lane 3BC29\.8±1\.829\.8\\pm 1\.831\.7±1\.731\.7\\pm 1\.71\.931\.930\.237±0\.0260\.237\\pm 0\.0260\.361±0\.0240\.361\\pm 0\.0240\.256±0\.0210\.256\\pm 0\.0210\.146±0\.0310\.146\\pm 0\.031AIRL28\.8±1\.828\.8\\pm 1\.830\.4±1\.730\.4\\pm 1\.71\.891\.890\.218±0\.0280\.218\\pm 0\.0280\.375±0\.0250\.375\\pm 0\.0250\.289±0\.0260\.289\\pm 0\.0260\.118±0\.0310\.118\\pm 0\.031GRID \(ours\)34\.7±0\.8\\mathbf\{34\.7\\pm 0\.8\}37\.9±0\.8\\mathbf\{37\.9\\pm 0\.8\}1\.980\.194±0\.0210\.194\\pm 0\.0210\.285±0\.0200\.285\\pm 0\.0200\.301±0\.0200\.301\\pm 0\.0200\.219±0\.0230\.219\\pm 0\.023

Table 2:Generalist Comparisons in the Highway\-Env\.MethodReward \(↑\\uparrow\)EpisodeLength \(↑\\uparrow\)Lane 3Absolute \(↑\\uparrow\)Lane OccupancyLane 0Lane 1Lane 2Lane 3Right Specialist22\.1±1\.722\.1\\pm 1\.723\.7±1\.723\.7\\pm 1\.70\.232±0\.0290\.232\\pm 0\.0290\.236±0\.0450\.236\\pm 0\.0450\.169±0\.0270\.169\\pm 0\.0270\.203±0\.0240\.203\\pm 0\.0240\.392±0\.0410\.392\\pm 0\.041Wrong Specialist34\.0±1\.634\.0\\pm 1\.634\.9±1\.534\.9\\pm 1\.50\.302±0\.0310\.302\\pm 0\.0310\.272±0\.0300\.272\\pm 0\.0300\.191±0\.0160\.191\\pm 0\.0160\.192±0\.0160\.192\\pm 0\.0160\.346±0\.0320\.346\\pm 0\.032BC29\.9±1\.829\.9\\pm 1\.832\.0±1\.732\.0\\pm 1\.70\.272±0\.0320\.272\\pm 0\.0320\.136±0\.0250\.136\\pm 0\.0250\.256±0\.0250\.256\\pm 0\.0250\.268±0\.0200\.268\\pm 0\.0200\.340±0\.0360\.340\\pm 0\.036AIRL24\.3±2\.024\.3\\pm 2\.026\.0±1\.926\.0\\pm 1\.90\.188±0\.0280\.188\\pm 0\.0280\.115±0\.0230\.115\\pm 0\.0230\.259±0\.0280\.259\\pm 0\.0280\.337±0\.0280\.337\\pm 0\.0280\.289±0\.0370\.289\\pm 0\.037GRID \(ours\)34\.9±1\.0\\mathbf\{34\.9\\pm 1\.0\}37\.9±1\.0\\mathbf\{37\.9\\pm 1\.0\}0\.349±0\.025\\mathbf\{0\.349\\pm 0\.025\}0\.139±0\.0250\.139\\pm 0\.0250\.218±0\.0200\.218\\pm 0\.0200\.274±0\.0140\.274\\pm 0\.0140\.368±0\.0240\.368\\pm 0\.024

Table 3:Specialist Comparisons in the Highway\-Env\.
## Appendix BExperiment Details and Hyperparameters

### B\.1Basis Function Decomposition Experiment

Data Generation\.For the synthetic basis function decomposition experiment, we generate 2000 demonstration samples uniformly distributed across the domain\[−4​π,4​π\]\[\-4\\pi,4\\pi\]\.

Network Architecture and Training\.The General Reward Inference and Disentanglement \(GRID\) model is trained to predict the total reward\. We use a sinusoidal positional encoding for the inputs\. The complete list of architecture details and training hyperparameters is provided in Table[4](https://arxiv.org/html/2606.18537#A2.T4)\.

HyperparameterValueGeneral network hidden size1024General network depth4Specific network hidden size4Specific network depth2Person embedding size \(dd\)2Embedding network hidden size128Embedding network depth2Emphasis loss weight \(λ1\\lambda\_\{1\}\)0\.05Embedding loss weight \(λ2\\lambda\_\{2\}\)0\.1Learning rate1×10−31\\times 10^\{\-3\}Epochs100Batch size128Table 4:GRID Architecture and Training Hyperparameters: Basis Function DecompositionAblation Study Details\.As discussed in the main text, the general reward is parameterized by a 2\-layer MLP with a hidden size of 1024, while the specific reward uses an intentionally bottlenecked 2\-layer MLP with a hidden size of 4\. To demonstrate the necessity of this capacity bottleneck, we conduct an ablation where the specific reward’s MLP hidden size is increased to 1024, matching the general reward network\.

### B\.2Craftax Experiment

Environment Setup\.To aid in visualizing GRID’s disentanglement, we constrain the original Craftax environment\. The modifications are as follows:

- •The environment map is reduced from a64×6464\\times 64grid to an8×88\\times 8grid\.
- •The action space for all agents is restricted tono\-op,up,down,left,right, anddo\-action\.
- •The original single\-step reward formulation is replaced with a cumulative reward structure\.
- •Adversarial entities \(e\.g\., skeletons, zombies, cows, arrows\) are removed\.
- •Agents no longer require tools to retrieve resources such as stone or diamonds\.
- •World generation randomization is largely disabled; only agent spawn positions remain randomized\.
- •Reward weights are modified to ensure each agent prioritizes a distinct resource\-gathering task\.
- •Episodes terminate immediately once an agent completes its specific task or is terminated\.

Data Generation\.We generate demonstration data using the IPPO implementation from\[[35](https://arxiv.org/html/2606.18537#bib.bib928)\]\. To capture a diverse range of skill levels, we collect a total of 2000 demonstration trajectories \(500 trajectories each from policy checkpoints at 1, 10, 100, and 200 training iterations\)\. The IPPO actor\-critic network utilizes a standard feed\-forward architecture comprising three hidden layers of sizes \[64, 64, 32\] with Tanh activations\. The IPPO training hyperparameters are detailed in Table[5](https://arxiv.org/html/2606.18537#A2.T5)\.

HyperparameterValueTraining iterations per agent200Parallel environments200Minibatches16Learning rate2\.5×10−42\.5\\times 10^\{\-4\}Entropy coefficient0\.02Max gradient norm1\.0Discount factor \(γ\\gamma\)0\.99GAE parameter \(λ\\lambda\)0\.95Clip coefficient0\.2Value function coefficient0\.5Auxiliary loss coefficient0\.1Table 5:IPPO Training Hyperparameters: Craftax Data GenerationGRID Training\.After collecting the IPPO demonstrations, we train the GRID architecture using the hyperparameters listed in Table[6](https://arxiv.org/html/2606.18537#A2.T6)\.

HyperparameterValueGeneral network hidden size1024General network depth4Specific network hidden size64Specific network depth4Person embedding size2Embedding network hidden size128Embedding network depth2Emphasis loss weight \(λ1\\lambda\_\{1\}\)0\.05Embedding loss weight \(λ2\\lambda\_\{2\}\)0\.1Learning rate1×10−41\\times 10^\{\-4\}Epochs100Batch size256Table 6:GRID Architecture and Training Hyperparameters: Craftax
### B\.3Highway\-Env Experiment

Environment Setup\.We adapt the standard single\-agent Highway\-Env into a multi\-agent environment featuring three agents navigating a four\-lane highway\. To define the specific behaviors of our heterogeneous population, Agent 0 is rewarded for occupying Lane 0, Agent 1 for Lane 1, and Agent 2 for Lane 2\. Lane 3 is left un\-preferred by all agents to serve as a holdout specialization target\. The environment’s reward coefficients are set as follows: the lane preference reward is 0\.1, the high\-speed reward multiplier is 1\.0, and the collision penalty is \-0\.5\.

Data Generation\.We modify the PPO implementation from Stable\-Baselines3\[[21](https://arxiv.org/html/2606.18537#bib.bib932)\]to support IPPO for our multi\-agent setup\. Table[7](https://arxiv.org/html/2606.18537#A2.T7)lists the hyperparameters used to train the expert demonstrators\. Unlisted parameters are default\.

HyperparameterValueLearning rate3×10−43\\times 10^\{\-4\}Steps per rollout \(n​\_​s​t​e​p​sn\\\_steps\)2048Batch size256Epochs10Discount factor \(γ\\gamma\)0\.95Entropy coefficient0\.05Table 7:IPPO Training Hyperparameters: Highway\-Env Data Generation![Refer to caption](https://arxiv.org/html/2606.18537v1/Figures/ippo_highway_lanes.png)Figure 10:IPPO Agents Lane Counts over 100 Episodes\. Average Episode Length=33\.033\.0GRID Training\.The GRID architecture is subsequently trained on the generated driving demonstrations to disentangle the universal safety/speed norms from the individual lane preferences\. The hyperparameters for this stage are provided in Table[8](https://arxiv.org/html/2606.18537#A2.T8)\.

HyperparameterValueGeneral network hidden size1024General network depth4Specific network hidden size64Specific network depth4Person embedding size1Embedding network hidden size128Embedding network depth2Emphasis loss weight \(λ1\\lambda\_\{1\}\)0\.05Embedding loss weight \(λ2\\lambda\_\{2\}\)0\.1Learning rate1×10−41\\times 10^\{\-4\}Epochs100Batch size256Table 8:GRID Architecture and Training Hyperparameters: Highway\-Env

## Appendix CBaseline Implementations and Hyperparameters

We evaluate GRID against Behavioral Cloning and Adversarial Inverse Reinforcement Learning baselines\[[21](https://arxiv.org/html/2606.18537#bib.bib932)\]\. Demonstration data come from IPPO’s expert policies\.

HyperparameterValueGeneral SetupExpert demonstrations1000Network hidden layers\[256, 256\]Behavioral Cloning \(BC\)Learning rate1×10−31\\times 10^\{\-3\}Batch size2048Epochs500Entropy weight0\.0L2 weight0\.0AIRL \(PPO Learner Policy\)Learning rate3×10−43\\times 10^\{\-4\}Steps per rollout \(n​\_​s​t​e​p​sn\\\_steps\)2048Batch size256Epochs \(n​\_​e​p​o​c​h​sn\\\_epochs\)10Discount factor \(γ\\gamma\)0\.95Entropy coefficient0\.05Table 9:Training Hyperparameters: BC and AIRL Baselines### C\.1KL\-Control Fine\-Tuning Hyperparameters

To evaluate the adaptability of the learned models, we fine\-tuned the BC, AIRL, and GRID priors to prefer a previously unseen target lane \(Lane 3\) using KL\-control\. Through empirical tuning, we found that Behavioral Cloning \(BC\) and Adversarial IRL \(AIRL\) required a stronger KL penalty \(β=0\.2\\beta=0\.2\) to maintain stability and prevent the policies from collapsing during adaptation\. In contrast, the GRID generalist prior was effectively fine\-tuned with a significantly lower penalty \(β=0\.05\\beta=0\.05\), indicating a more robust foundation that easily accommodates new objectives\. We additionally evaluated specialized baselines, denoted as the Right Specialist and Wrong Specialist\. The optimal learning rates and KL penalty coefficients for all methods during this fine\-tuning phase are detailed in Table[10](https://arxiv.org/html/2606.18537#A3.T10)\.

MethodKL Control Coefficient \(β\\beta\)Learning RateBC0\.201×10−31\\times 10^\{\-3\}AIRL0\.201×10−51\\times 10^\{\-5\}GRID Generalist \(Ours\)0\.055×10−55\\times 10^\{\-5\}Wrong Specialist0\.055×10−35\\times 10^\{\-3\}Right Specialist–5×10−45\\times 10^\{\-4\}Table 10:KL\-Control Fine\-Tuning Hyperparameters

## Appendix DComputation Resources

Basis Function decomposition experiments were performed on NVIDIA Ryzen 9 7900X and GeForce RTX 4080 with 32GB RAM\. Craftax experiments were performed partly on an internal cluster and on Ryzen 9 7900X and GeForce RTX 4080 with 32GB RAM\. Highway\-Env experiments were performed on i9\-14900K and NVIDIA GeForce RTX 5090 with 96GB RAM\. Runtimes were not recorded, but were not out of the ordinary\.