Heterogeneous Knowledge for Augmented Modular Reinforcement Learning
Authors: Lorenz Wolf
Mirco Musolesi
Venue: Accepted by TMLR ()

Abstract: Existing modular Reinforcement Learning (RL) architectures are generally based on reusable components, also allowing for ``plug-and-play'' integration. However, these modules are homogeneous in nature - in fact, they essentially provide policies obtained via RL through the maximization of individual reward functions. Consequently, such solutions still lack the ability to integrate and process multiple types of information (i.e., heterogeneous knowledge representations), such as rules, sub-goals, and skills from various sources. In this paper, we discuss several practical examples of heterogeneous knowledge and propose Augmented Modular Reinforcement Learning (AMRL) to address these limitations. Our framework uses a selector to combine heterogeneous modules and seamlessly incorporate different types of knowledge representations and processing mechanisms. Our results demonstrate the performance and efficiency improvements, also in terms of generalization, which can be achieved by augmenting traditional modular RL with heterogeneous knowledge sources and processing mechanisms. Finally, we examine the safety, robustness, and interpretability issues stemming from the introduction of knowledge heterogeneity.

Published in Transactions on Machine Learning Research (11/2025)
Heterogeneous Knowledge for Augmented Modular Rein- forcement Learning
Lorenz Wolf
lorenz. wolf 22 22 uk
Deportment of Computer Science Centre for Artificial Intelligence University College London
Mirco Musolesi
m. ww..usoleesiuccllac..o
Department of Computer Science & Centre for Artificial Intelligence University College London
Department of Computer Science and Engineering Universily of Bologna
Reviewed on OpenReview: https: openreview. forum? id= 2me87Ybs.0D
Abstract
Existing modular Reinforcement Learning (RL) archilectures are generally based reusable components, also allowing for "plug-and-play" integration. However, these modules are homogencous nature. In fact, they essentially provide policies obtained via RL through the maximization of individual reward functions. Consequently, such solutions still lack the ability to integrate and process multiple types of information (i.e., heterogeneous knowledge representations), such as rules, sub-goals, and skills from various sources. In this paper, discuss several practical examples of heterogeneous knowledge and propose Augmented Modular Reinforcement Learning (AMRI.) to address these limitations. Our framework uses selector to combine heterogencous modules and scamlessly incorporate different types of knowledge representations and processing mechanisms.. Our results demonstrate the performance and efficiency improvements, also in terms of generalization, which can be achieved by augmenting traditional modular RL with heterogencous knowledge sources and processing mechanismss. Finally, examine the safety, robustness, and interpretability issues stemming from the introduction of knowledge heterogeneity.
1 Introduction
Reinforccment learning (RL) has emerged powerful paradigm for sequential decision-making in complex and dynamic environments (Sutton & Barto 2018 Kiran et al. 2022 Its successes extend across wide range of domains, including robotics (Kober et al. 2013) dynamic resource allocation (Waschneck et al. 2018), and game playing (Mnih et al. 2015 Silver et al. 2016 Vinyals et al. 2019). Despite these advances, the practical deployment of RL in real-world systems remains constrained by two fundamental challenges: low sample efficiency and safety concerns arising from exploratory behaviors that may lead to high-risk actions in critical scttings (García & Fernández 2015).
One approach mitigating sample inefficiency is modular RL, in which tasks decomposed into sub- problems and solved by specialized skill modules. These modules can be reused across contexts, accelerating
learning and improving flcxibility (Jacobs et al. 1991 Russell & Zimdars, 2003 Spraguc & Ballard 2003; Simpkins & Isbell 2019 Devin et al. 2017). Prior modular and hierarchical RL (HRL) frameworks typically assume homogeneous knowledge among modules. That is, most sub-policies are trained via RL under com- parable reward structures. This reliance on homogencity presents both thcoretical and practical barricrs incorporating diverse knowledge sources. For example, rule-based controller, retrieval mechanism, and RI.-triined policy may differ not only in their representations and processing pipelines but also in their

---

Published in Transactions on Machine Learning Research (11/2025)
Update Knowledge
Environment
Update Knowledge
Action Proferences
AMRL
Action Preferences
Module
N-1
Modila
Rules for Safety: "Lava is not safe"
Retrieval Systemm: from trajectory dataset
Subgoals and instructions for eff efficiency
Skill: Ri--asedd policy
Heterogeneous Knowledge
Figure/ 1: Examples of heterogeneous knowledge and the AMRL architecture. AMRL is able to access several sources of heterogencous knowledge via modules. The modules then be updated based the environment feedback.
action-value semantics, making direct integration challenging without architectural modifications. Existing frameworks such as KoGuN (Zhang et al. 2020) and KIAN (Chiu al.) 2023) illustrate this limitation: while modular in structure, they are not explicitly designed to accommodate modules with fundamentally different internal representations, leading to performance degradation when heterogeneity is introduced.
However, in practice, diverse knowledge sources are often readily available and highly complementary. Rules and logical constraints can encode safety-critica requirements; pre-trained skills can provide strong priors; retrieval mechanisms can recall past scenarios; and RT. modules can adapt online to optimize performance. Leveraging this heterogeneous knowledge, i.e,, information derived from distinct paradigms, each with its own representations and processing mechanisms, has the potential to enhance sample efficiency, safety, and robustness simultaneously. While previous work has explored integrating single alternative sources, such text manuals (Zhong ct al. 2020): demonstrations (Abbeel et al., 2010), knowledge graphs (Murugesan et al. 2021), these efforts remain constrained by their focus on one type of knowledge at time. There is currently no general framework that flexibly and systematically integrates diverse sources of heterogeneous knowledge.
To address this gap, we propose Augmented Modular Reinforcement Learning (AMRL (Figure uni- fying framework that is explicitly designed support the seamless integration of heterogeneous knowledge modules. In AMRL, all modules consume the same environment state and act within the same action space, but critically, they process the state differently: rule-based modules extract abstract features (e.g., "is there lava ahead?"), retrieval modules rely on stored experiences, and skill modules employ task-specific learned representations. This ensures that heterogeneity is not superficial but reflects genuinely distinct internal interpretations of the environment. AMRL introduces a general selector that arbitrates among modules through two mechanisms: hard selection, corresponding to command arbitration (Simpkins & Isbell 2019) and soft selection, akin to command fusion (Russell & Zimdars, 2003) and also certuure-of-xppert. (Shazoer let al. 2017).. While command arbitration and command fusion are not novel frameworks in isolation, we de- play these selection mechanisms incorporate heterogeneous knowledge and show that their behaviors differ significantly in the presence of heterogeneous modules: soft selection, in particular, improves robustness and mitigates the negative impact of noisy or uninformative modules.
As a motivating example, consider an autonomous traffic control system at a busy intersection. An effective system would need to: (1) follow traffic laws expressed as rules, (2) reuse pre-trained policies from similar interscctions, (3) optimize flow dynamically through RL, and (4) retrieve past cases handle emergencies. Such system inherently requires heterogeneous modules with distinct representations and reasoning mech- anisms. AMRL enables precisely this type of integration, providing principled architecture that abstracts away representational differences while ensuring safe and efficient operation.
The full implementation and code used for experiments are publicly available:
2

---

Published in Transactions on Machine Learning Research (11/2025)
In summary, we make the following main contributions:
Unified formalization of helerogeneous modules for decision-making We introduce framework for representing and integrating diverse knowledge sources rules, logic, skills, and retrievals -within a modular RL system.
Augmentation of modular RT. We extend classical modular RT by enabling plug-and-play integration of heterogencous modules through a selector mechanism that is agnostic internal representations, Ihereby improving sample elficiency, safely, and robustness.
Comprehensive evaluation of heterogeneous knowledge integration. We empirically demonstrate that AMRL not only improves sample efficiency but also enhances robustness, safcty, generalizability, and interpretability. This evaluation offers new insights into the benefits and trade-offs of modular approaches when enriched with external knowledge.
In this regard, AMRL represents substantive augmentation of modular RL, offering an empiricallly-aalidated framework that explicitly and systematically integrates heterogeneous knowledge.
2 Related Work
Knowledge augmented RL. Grounding RL agents with external knowledge is an ongoing line of rescarch in the field of RL. For example, progress has been made on incorporating knowledge provided to the agent via text for example as manual for a game (Zhong et al. 2020 or knowledge graphs connecting concepts and their characteristics (Murugesan et al.] 2021). One approach to incorporate knowledge sources is the formalism of contextual MDP (CMDP) (Kirk et al., 2023 Perez et al., 2020 Ghosh et al. 2021 Hallak et al.) 2015).
Other approaches equip agents with retrieval mechanisms access memory of past trajectories to inform
current decisions [Goyal et al. 2022 Humphreys et al. 2022).
While these methods demonstrate the value of augmenting RL with external knowledge, they remain limited in scope, integrating only single type of knowledge (e.g., trajectories, rules, or symbolic abstractions). They typically assume that knowledge can be processed within a uniform framework. Hybrid approaches, such as the hierarchical planning method by Lorang et al. (2024) and symbolic recovery in Goel el. al. (2022). partially address broader knowledge integration but are domain-specific. Modular frameworks like KoGuN (Zhang al., 2020) and KIAN (Chiu et al.. 2022 2023) support multiple modules, but they still assume relative homogeneity, either in representation AMRL explicitly formalizes the integration of heterogeneous knowledge sources, including rules, retrieval modules, and pre-trained skills, through a unified interface that is agnostic to differences in representation and processing, cnabling flexible composition.
in how modules trained and combined. In contrast,
Hierarchical and Modular RL. Hicrarchical RL (HRL) decomposes problem into several smaller subtasks organized in hierarchy (Digney 1998). The higher-level parent-tasks can call lower-level child- tasks as if they were primitive actions. This enables higher-level tasks to focus on more abstract and longer-term learning while lower-level tasks are responsible for primitive actions and fine-grained control (Hengst 2010). Each task is learned by a policy and at decision time; the higher-leve policy triggers one in the level below. The goal is for the agent to perform task abstraction by focusing on more high-level tasks, that it will suffer less from the curse of dimensionality (Barto & Mahadevan, 2003). This is a very active area of research with some of the earlier works dating back to the 1990s (Dayan & Hinton 1992. Vezhnevets et al. 2017 Parr & Russell 1997). HRL is closely related to modular RL, which aims to decompose the learning process into complementary policies/
Dalicies/modules modules that can be combined by controller (Sutton et al. 2011 Goyal el al. 2020: Jacobs et al.) 1991). Previous works in Modular RL have considered command arbitration, selecting one module, which takes over the agent for the next action, and command fusion, combining the preferences of the different modules (Gupta et al.) 2202 Russell Zimdars, 2003 Mendez Andreas et al. 2017). We leverage the same modularity and flexibility for heterogeneous knowledge.
let al. 2022
3

---

Published in Transactions on Machine Learning Research (11/2025)
Command arbitration and fusion. The selector used by AMRL is closely linked to the arcas of command arbitration and command fusion. In the command arbitration paradigm an arbit decides which module takes control over the agent in a given state; the selected module then executes an action. Earlier works include GM-Sarsa (Sprague & Ballard, 2003) and Q-decomposition (Russell & Zimdars 2003). More recent methods Arbi-Q (Simpkins Isbell| 2019) and GRACTAS (Gupta et al. 2021) assume that the arbor receives a global reward signal to be maximized, enabling it to handle different scales of rewards for each of the modules. Furthermore, the global reward signal means that these methods can be used with off-policy as well as on-policy learning (Simpkins & Isbell, 2019). The underlying mechanism in command fusion is also akin to the gating function in Mixture-of-Experts (MoF) architectures (Shazeer et al. 2017 Chen et al. 2022), which has been widely used in Deep Learning and also in Large Language Models. These prior works focus on the arbitration and fusion of homogemeous RL policies. In contrast, AMRL investigates these mechanisms in the context of heterogeneous modules, where differences in representation, processing, and reliability introduce new challenges. Our experiments show that soft selection (fusion) can learn to miligate conflicts and noise among diverse modules, improving robustness, while hard selection (arbitration) may amplify instability when modules provide conflicting or uninformative guidance. This highlights that classical selection strategies must be re-evaluated when extended to heterogeneous integration.
3 Preliminaries
Let .M- = (8, A,P, R, pa, 7, T) represent discrete-time finite- horizon discounted Markov decision process (MDP). .M consists of the following elements: the state space S, the action space A, transition probability distribution P: SSSx x A x S - R2, reward function R S x A + R, an initial state distribution Pa : S - R+. a discount factor 7 E [0, 1J, and the horizon T. Generally, the policy T² : S x A + R_ is parameterized by 8, which we optimize in order to maximize the expected discounted return under the policy. We denote the probability of taking action a in state St by (a|s)) and abuse notation by denoting the distribution over all actions with x xst). Even if the architecture concoptually independent from the underlying RL algorithm, the design and inplementation presenled in this paper relies on Proximal Policy Optimization (PPO) (Schulman et al. 2017) for discrete action spaces and Soft Actor-Critic (SAC) (Haarnoja et al. 2018 in the continuous case.
Moveover, in many real-world deployments, the underlying state is partially observable. Formally, these scenarios can be modeled as a partially-observable Markov Decision Processes (POMDPs), which is repre- sented by the tuple AS S, O, R, Po. 7, T), where in addition to the already known quantities we have introduced the observation space f and the observation probability distribution O()). Instead of s S, the agent now observes Ω generated from the state via o ~ 0(8).
4 Heterogeneous Knowledge Sources, Representation, and Processing
Heterogencity can arise for various reasons. As a consequence, AMRL agents not only rely on heteroge- neous knowledge sources but also on heterogeneous representation and processing mechanisms (Table D. Knowledge from the same source can be represented in various ways and, naturally, rules must be processed differently to trajectories. Importantly, this goes beyond prior modular hierarchical RL frameworks, which typically assume homogeneous modules consisting of RL-trained sub-policies with comparable reward scales, and therefore do not consider modules with fundamentally different representational formats or reasoning paradigms. We will adopt a bottom-up approach in the description of AMRL: we first discuss the modules at the basis of our approach; they essentially act as containers for heterogeneous knowledge sources and pro- cessing. In particular, we focus on the four types of modules outlined in the motivating example, specifically, rules, skills, retrieval, and dynamic RL. We then discuss how these modules are incorporated in the overall architecture via selection mechanism based on hierarchical modular RL.
4.1 Knowledge Representation and Processing using Modules
An AMRL agent instantiated with several modules M, E M, which each act as a container for knowledge representation K2 and a corresponding processing mechanism. Module M, generates action preferences,
4

---

Published in Transactions on Machine Learning Research (11/2025)
Table 1: Knowledge heterogeneity can stem from various factors, namely the source of the knowledge, its representation, and the processing methods applied.
Source
Representation
Processing
Human Expert
Rules
Logic
Environment
RL-based policy
Skill Execution
Previous Deployment
Trajectory Database
Retrieval
represented as probabilities over possible actions, based on its knowledge source. This is expressed mathe- matically as TM, TM, (a|s;) = MM, st/K). In contrast to previous works in modular RL, the AMRL modules are heterogencous. They are not limited to RL policies and can be diverse in terms of how they form preferences over actions. Modules and their knowledge sources can be either static or dynamic. If the knowledge source, and consequently the mapping of module Mi, needs be updated during training, modular feedback must be observed. This ensures that the effects of actions proposed by M; can be learned effectively. Modular rewards {ra, J where ri, is the reward for module M, at time t, can be specified to favor modular learning. In case M; is static no modular updates are performed.
We now focus on the knowledge K; contained in modules M. Knowledge sources can be represented in many different forms, inchuding rules, manuals, trajectories, etc. and can be assessed according to the following characteristics: actionability, interpretability, and informativeness. A knowledge representation is actionable if it can be used directly inform the action selection process. Most actionable choices are procedural declarative knowledge sources, for example in form of world model (Hafner et al. 2025) or database of past trajectories (Goyal et al. 2022 Humphreys et al. 2022). While these might be the most effective reprosentations for RL agents, they are less interpretable and importantly can be harder to obtain in the first place due to deployment restrictions. Other more abstract knowledge representations such as rules and manuals are more interpretable, easier to validate and in many cases more widely available, which highlights the importance of flexibly incorporating and acquiring heterogencous knowledge. The informativeness cap- tures how relevant and useful the knowledge is for solving the given task. In our evaluation we compare the effect of these properties on the agent's performance and sample efficiency.
4.2 Types of Modules
In this subsection, we provide practical examples of heterogencous modules. It important to note that the proposed architecture flexible and supports multiple types of modules, each based on different knowledge representations and processing mechanisms. In the following, we select modules that are representative of broader classes of solutions. We first describe each module type for discrete action spaces and then outline how it can be exlended to continuous action spaces. We are also aware of ongoing research on each of these [Liu et al., 2023: Likmeta et al. 2020 Hasanbeig et al., 2023 Goyal et al. 2022 Humphreys et al. 2022), but note that our focus is on combining heterogeneous decision makinn mechanisms. A more detailed examination of each possible mechanism is beyond the scope of this paper.
4.2.1 Logic-based Rules
This module relies on a set of rules A = ( where each rule (;(c;) formalized as statement of the form:
If 9t satisfies condition o then TM (alst) (a)(e)(a).
where Px(a)) (a) is the probability of taking action under rule X; with condition a satisfied. For example, let us consider one of the typical benchmark games where an agent has to move in space where a lava flow is present: if there is lava in front of the agent then the probability of moving forward is set to 0. Conversely, let us suppose to have a different game environment in which, for instance, an agent has to collect objects, such as keys: if there a key in front of the agent then the probability of picking up the key is set to Rules can be organized hicrarchically with one rule potentially relying on other rules. Furthermore, rules contradict coincide in which case the conflict needs to be resolved. As this is not the main focus of
5

---

Published in Transactions on Machine Learning Research (11/2025)
this work, simply resolve such conflicts by averaging the probabilities over the set of rules applicable in state & denoted by Ar, such that:
TM (a/ss, A)
Σ (a())(a)
(1)
[Ael
AEA.
4.2.2 Trajectory Database with Retrieval
This module relies database containing trajectories and retrieves relevant information via nearest- neighbor search. The knowledge source in this case is given by the set D { where each (18.1,4;11,...)
(st,7, O,, Ti,2;)) )) is trajectory of length T; containing state-action-rewarr tuples. The retrieval mechanism relies on an embedding network mapping from the state space to an embedding space. Given the current state the query formed via qi = embed((s)) st). The k approximate nearest- neighbrrs in the embedding space are retricved from D with the Facebook AI similarity search library (FAISS) (Johnson al. 2021 based on the L2-norm, which yields the k tuples (Sx1: (m11:an1+ :..)...., (S One Tne) with corresponding distances day....,, dng to the query- The retrieved information is utilized to form action by weighting according to rewards and distance, which yields action preferences:
(PD(a)
Σ Tny I(a ang)
(a D)
EaEA (Po(a) with PD(a)
(2)
dny
We note that other more sophisticated methods for retrieval such as those proposed by Goyal et al. (2022)
and Humphreys et al. (2022 can be substituted instead.
4.2.3 RL-based Policies (Skills)
The third class of modules we consider is that based skills, i.c., on policies Takitt trained via RL with unknown reward function. The policies are trained on the same state space S with action space A A. These modules simply contain the policy and return corresponding action probabilities given by TM (alst) = (a)s) if a E A& and 0, otherwise.
4.3 Extending Modules to Continuous Action Spaces
For continuous state and action spaces, the conditions may be expressed inequalities ranges, and instead of discrete probabilities, the rule produces a parametric distribution over actions, e.g., A)= ((Ha($)), (s:)) (s:)) where (s(s) and (st) can be constants or functions determined by the rule. The retrieval module can be adapted for contimuous action spaces by replacing the indicator function with a kernel density estimate over the retrieved actions, yielding smooth probability distribution. Additionally one needs replace the softmax over discrete actions with normalization of the kernel density. Skill modules in the continuous action space return parameters of distribution, e.g., Gaussian N (usamm ( instead of discrete probabilities.
5 Combining Heterogeneous Knowledge
To incorporate various types of knowledge sources and processing mechanisms as described in Section 4 we employ a controller / arbitrator, which we refer to as the selector. The goal of this component is to select/combine one or more modules. In particular, it determines which modules to trigger in a given state, thereby enabling the performance of tasks with varying levels of complexity.
We consider two variants of selection mechanisms, which correspond to command arbitration (Sprague & Ballard 2003) and command fusion [Russell & Zimdars 2003). The first, which we refer to as hard selection, restricts the weights to onc-hot vector so that exactly one module is selected and executed at cach time step. The second, which we refer as soft selection, relies on an attention mechanism form weighted average combining preferences from several modules, resulting in command fusion. Thus, by weighting the

---

Published in Transactions on Machine Learning Research (11/2025)
modules, this variant of the selector is based in sense mixture model. We denote the selector's policy é mapping from the state space to the selector's action space by "selector S - Axelector, parameterized by Φ. We have Aselector f.i..... My). such that T (M: is the probability of choosing module M, in state &t and expresses the selector's preference for module M,- In the following, will provide formalization of the two selection mechanisms and discuss how to bridge the gap between the current task and what is possible by flexibly combining the modules.
5.1 Hard Selection
At each time step 4, the selector observes the current state and selects exactly one module M to execute. We then sample an action from the chosen module a, <M())). In particular, one module is sampled according to the selector's policy M Th Teelector To perform hard selection in differentiable manner, generale sample vector with dimension |Aselector N by setting:
exp ((log (
TT seleector (Teeeecoor + 9: /T)
V1
(3)
N Σ
Lj=1 exp ((lo
T solectoor (M]ls) - g3 /T)
for i- 1,..,, N, where 91 ... -N are i.i.d samples drawn from Gumbel 1), which can be sampled using an inverse transform (Jang et al. 2017]. The obtained sample is subsequently discretized into one-hot vector. This ensures the selection of exactly one module at a time. soft approximation is used for the computation of the gradients. This yields as the AMRL agent's policy =(ass.) = т (als)) with M n Taelector (8t). In the experiments we the temperature 1. We provide an ablation analysis in Appendix F.
5.2 Soft Selection
In contrast to hard selection, soft selection combines the modules' policies with a weighted average. The action preferences TM, (als;) returned by the modules are weighted by the module preferences Teteetor (Milss) returned by the selector. In particular, we have that r(as;)) = E @ seelector (M,ss (als:). Note that "setector (8)) is normalized to to 1.
5.3 Implementing the Selectors
In many applications, ftexibly combining the knowledge available the agent via the modules may not be sufficient to achieve high performance
task. Similarly to the inner actor in the KTAN architecture, ~ bridge this gap by adding dynamic RL module. This module contains learnable RL policy Mays
S x A + R4 trained to maximize the global reward function also observed by the selector. In practice, the selector implemented as a neural network parameterized by parameters ø, which, given the current state as input, outputs probabilities over the modules. By including the dynamic RL module, we have as AMRL policy with soft selection:
0.0 (a = = selector (Magn "d (als) +
Σ
T selecctr (M
(4)
MEM
which allows us to easily compute policy gradients w.r.t. Φ and 0, justifying the use of PPO to train the selector (see Appendix A for details). Note again that the selector itself is agnostic with respect to the actual knowledge representation. The entire decision-making process is presented in Algorithm
5.4 Extending AMRL to Continuous Action Spaces
We now discuss the extension of the AMRL framework to continuous action spaces. The hard selection mechanism can easily applied to continuous-aco acttion modules by selecting a single module's Gaussian policy at each time step. The action is then sampled from the selected distribution, rather than a discrete action. Gradients can be computed using the Humbel-Softman approximation over the module selection weights.

---

Published in Transactions on Machine Learning Research (11/2025)
Algorithm Decision-making with an AMRL agent in discrete action spaces.
Require:
Modules and knowledge M {M, k. Mayn
50 Initial state
for each timestep = O,... do
Get module preferences TM; (8t) t (i)
if Hard selection then
Sample selected module M ~ क setector (8c)
Get AMRL policy x{a(st) = Tas (c)st)
else if Soft selection then
T[a st) MEM TI selecctor (M st (a(ss)
end if
Sample action from AMRL policy a (8))
Observe St+1 and reward
Update selector parameters & and dynamic module parameters with PPO end for
Soft selection in contimouss-action domains can be performed by forming the true mixture of Gaussian distributions for the modules. However, this would require computing the full weighted sum of densities, which makes sampling and entropy computation intractable and can lead to high-variance gradients. Instead, we approximate the fused policy as single Gaussian whose mean and variance are the weighted sums of the individual modules' means and variances:
(sused($)) Σ wipp(se)s o
2 iueed = E (s))
(5)
where wi are the selector weights and (V (f()) are the individual module distributions. This approx- imation enables efficient sampling and stable gradient computation using the reparameterization trick, while retaining contributions from all modules. Forming a true mixture of Gaussians would require computing the logarithm of the sum exponentials (log-sum-ezp) over all components, which can increase gradient vari- ance and reduce learning slability, particularly when the number of helerogeneous, partially deterministic knowledge modules large. We adopt this single- GaulsGan approximation to ensure robust and tractable learning with heterogencous knowledge sources in continuous-action domains.
5.5 Module Prioritization
With heterogeneity come inherent differences between the modules. In particular modules such as the rulc module may be considered safe and worstworthy, carefully designed and well understood. Consequently, while the rules may not be applicable in all cases, if applicable, they can be trusted and should be followed by the agent. This can be achieved with module prioritization. In the case of the rule module, whenever triggered, the rules act as constraint, i.e. other modules' action preferences can be
more of its rules taken into account, but only long as the agent does not violate any of the triggered rules. This can also be seen as masking of actions. In other words, for a given state 8, the actions, which are not to be taken, are extracted from, for instance, the rule module in form of a mask defined by Ag = if a an unsafe action and otherwise. Then the final policy is given by & where n (als) is defined safety critical states, the applicable rules mask unsafe actions, which are not be taken. This reduces the percentage of unsafe actions 0 given that the rules cover all cases. Similarly to the soft and hard selection mechanisms, soft prioritization mechanism based on specified weights can also be implemented.
Equation 4 Tn
6 Evaluation
In the evaluation of AMRL (and solutions for the integration of heterogencous knowledge in RL more gener- ally) we focus on the overall performance and sample efficiency, the dependence on knowledge informativeness and robustness against random modules, the safety gains during training and at test time, and the inter-
8

---

Published in Transactions on Machine Learning Research (11/2025)
pretability of the selection mechanism. Unless specified otherwise explicitly, we use the acronym AMRL to denote AMRL with the soft selection mechanism.
6.1
Experimental Settings and Baselines
6.1.1
Overview
We now evaluate AMRL considering a range of modules across different. benchmark environments. In par- ticular, we train the agents on both discrete and continuous action spaces.
Discrete Action Benchmarks (Minigrid). We use several environments from the Minigrid suite Chevalier-Boisvert al. 2018), each presenting distinct challenges:
DoorKey: The agent must locate and pick a key, unlock a door, and reach a goal square. Evaluated on 6x6 and 8x8 grids.
CrvaCiossigg: A safety-critical environment where the agent navigates through streams of lava to reach a goal. Stepping on lava terminates the episode. Evaluated on 9x9 grids with 1 or 2 lava streams.
Empty: A simple navigation task where the agent moves goal square without obstacles. Evalu- ated 16x16 grid.
In all Minigrid environments, the agent has 7 discreto actions and observes a 7x7 grid contered on its position. Additional details on action space and visual representations are provided Appendix
Continuous Action Benchmark (OpenAL-Roboties). For evaluation continuous action spaces, we use the Fetch enviromments (Plappert et al. 2018), set of manipulation tasks performed with a 7-DoF robot from the OpenAl Robotics Gym qe Lazcano al. 2024). Tasks include reaching, pushing, sliding, and pick-and-place, which require precise control over continuous joint velocities and gripper actions.
Modules Evaluated Across Benchmarks. We consider instances of the module types described in Sec- tion 4.2 For Minigrid, the rules module (Section 4.2.1 contains a set of 6 rules relevant to the environments enforcing constraints such as "do not step onto lava" (Figure 2 Violations result in probability zero being assigned to unsafe actions. The retrieval module (Section 4.2.2 relies on expert trajectories. Lastly, the RL-based skill module (Section 4.2.3) leverages pre-trained policies specific the skill environment. For the Fetch environment on the other hand, the first module moves the gripper closer the object, while the second module contains a rule-based policy to approach the target position.
While all modules receive the same raw environment observation, the way they interpret and process this input differs substantially. Each module applies its own internal transformations and abstractions: the rule- based module distills the state into interpretable features like "is there lava ahead?", whereas the skill and retrieval modules employ learned state representations optimized for their specific tasks. This heterogeneity in internal processing allows the system to combine multiple perspectives on the environment, leveraging both structured, human-designed rules and flexible, learned behaviors, even though the external observation interface remains unified. All experiments are conducted with 10 random seeds"] and we report mean and sample standard deviation.
6.1.2 Baselines
As baselines, we implement KIAN Chiu et al., 2023 KoGuN ZZang et al. 2020), and standard rein- forcement learning algorithms: PPO Schulman et al.] 2017) for discrete action spaces, and SAC (Haarnojal et al. 2018) for cont inuous action spaces. KIAN and KoGuN are two cutting-edge methods for knowledge- augmented RL and are adapted to our scenario, which involves heterogeneous knowledge, to make them as competitive as possible. The implementations of all agents with discrete action spaces rely on the
Experiments on the robotics environment based only serds at the moment, the remaining seeds being collected.
9

---

Published in Transactions on Machine Learning Research (11/2025)
there is key inside the state, then pick it up if possible or clsc towards it.
there is ball inside the state, then pick up if possible or else move towards it.
If there is door inside the state, then open it if possible or clse move towards it.
there is goal inside the state, then
towards
there is Lava in front the agent, do not move forward.
If there is wall in front of the agent, do not move forward.
Figure 2: Set of rules used the MiniGrid evaluation.
rl-starter-files repository and torch_ac a standard settings from rl ti-baselines3-zn [Raffin] 2020) and mabllebaaseiness (Raffin et al. (021).
For continuous action spaces, SAC is implemented following
Default hyperparameter scttings are used for both PPO and SAC. Additional implementation dctails are provided Appendix C
KIAN learns key for each piece of knowledge (originally rules) via an embedding layer. The actor consists of inner component, which plays the same role the dynamic RL module in our architecture, query network that, given the current state, forms query used to weight the knowledge pieces, and another fully connected layer to learn the keys of each knowledge piece. We have directly adopted the official implementation.
KoGuN, implemented with heterogeneous knowledge, evaluates each module and concatenates the outputs form k(s,) (MM, (se); ..., M, (st)). The knowledge vector h(s;) is then fod into the actor network part of the state representation, such that (a]s:) = Te (k(ss))), st).
PPO and SAC serve as our standard RL baselines that do not incorporate additional knowledge.
6.1.3 Computational Cost
The computational cost. for training and deploying AMRL is comparable with that of KoGuN. KIAN has a larger number of trainable parameters and a slightly slower inference time. PPO is faster at inference time since does not rely on any of the knowledge modules. The details of this analysis can be found Appendix D
6.2 Heterogeneous Knowledge for Efficiency Across Environments
To evaluate how heterogeneous knowledge improves learning efficiency across range of environments all agents are equipped with the following knowledge. In particular, we consider Skill, which solves the Minigrid Unlock environment; Retrieval, which is based on trajectory data collected in Empty Random 5x5 by an agent trained on the Empty Random 5x5 task; and, finally, Rules, which relies on all 6 rules for Minigrid. AMRL and baselines benefit from heterogeneous knowledge. In particular, AMRL with soft selection is more efficient and achieves higher performance than baselines across most environments (see Table 2). Not only does AMRL achieve strong final performance, but Figure also shows its improved sample efficiency across several Minigrid environments. Notably, the performance difference is largest on the DoorKey 8x8 environment for which other methods have not managed to make meaningful progress by the time AMRI. has solved the environment.
While AMRL with soft selection performs well, the hard selection mechanism leads "noisy" behavior and poorer performance, often with slower convergence and without reaching competitive final performance. This should be expected due to its inability of combining the action preferences from several modules, but it is also dependent on the temperature parameter т (see Appendix F.3). We note that all three modules are highly informative for solving the DoorKey environments in contrast to only 2 modules for Empty (Retrieval & Rules) and only module for Lava Crossing (Rules). This is apparent obsorving the relative performance of agents accessing knowledge against PPO, which improves as the module informaliveness decreases.
TThe code of the libraries can found at the following URLs: httpe://ggthab..cmattte//il http:///...ccceeeeeeee///eeeeeaa https://githun htt.:///iiieeee//rrrreaaaaa
and
10

---

Published in Transactions on Machine Learning Research (11/2025)
DoarKey
DoorKey 8x8
LavaCrossing S9N1
Emply 16x16
Rewardd
Rewardd
Reward
WWRL_hard
KoGuN
Frames
Frames
Hrames
Frames
Figure 3: The achieved reward logged Uhroughout training. On all 4 environments AMRL with soft selection uses heterogeneous knowledge to achieve good performance more efficiently than baselines. The hard selection mechanism strongly limits its capabilities and results in significantly nosier behaviors.
Unsafe actions
Weight of module
0.1
Dynamic
0.88
Skill
0.08
Retrieval
Rules
0.04
%
0.02
0.00
AMBB
KIAN
KrCaN AMTS, Jard
PPO
Agent
Frames
(a) Safety.
(b) Interpretability.
Figure 4: 4aA Percentage of unsafe actions during training on LavaCrossing S9N1. Confidence intervals are + standard deviations across random seeds. AMRL performs the smallest amount of unsafe actions closely followed by KTAN. 4bb AMRL selector weights during one episode evaluation on DoorKey 8x8 showing mean #2 standard deviation calculated across 10 training seeds. On average, agents pick up the key with their 6th action and unlock the door with their 11th action, indicated by the dashed lines. The second milestone causes shift in the module weights away from the unlock skill and Lowards the dynamic module.
6.3 Safer Training and Deployment
A strong motivation for utilizing heterogeneous knowledge in RL is to improve safety during training and deployment. All agents incorporating heterogeneous knowledge achieve significantly higher performance than PPO on the safety critical Lava Crossing environments, which directly corresponds to fewer unsafe actions. We record unsafe actions, such as crashing into a wall or coming in contact with the lava, during training on the LavaCrossing S9N1 enviromment; we plot the percentage of unsafe actions performed by each agent in Figure AMRL performs the smallest amount unsafe actions averaged across training runs.
Beyond the above safety improvements achieved by incorporating hetcrogencous knowledge, module priori- tization can be deployed with AMRL for further gains. Given that module is known to produce safe and trustworthy actions only, such as the rule module, module prioritization can be deployed to shorteut the decision making process and only act according to the safe module in critical states. This will reduce unsafe actions to 0% if the module is perfectly accurate.
6.4 Improving Interpretability Through Modularity
The modularity of AMRL combined with heterogeneous knowledge significantly improves interpretability. Heterogeneous knowledge sources, such rules and skills, are inherently more interpretable than neural network-based policy trained on the fly. Assuming interpretable modules, an analysis of the selector's 's outputs is sufficient to gain an insight into performed actions. Additionally, we note that the number of modules is likely to be smaller than that of the primitive actions, reducing the complexity of an analysis. Plotting the selector's weights during an evaluation episode in DoorKey (Figure 4b) shows the expected order of first relying on the Unlock skill and then the dynamic module. Additional results for other environments and agents can be found in Appendix F.2
11

---

Published in Transactions on Machine Learning Research (11/2025)
Table Evaluation of heterogeneous knowledge for generalization and efficiency across environments. Agents use the following heterogeneous knowlodge: skill="Unlock" retrieval= "Empty rondio__nndd__bb__ 5",
rules="all! AMRL with soft selection outperforms except for the Lava, S9N2 environment (note large standard deviations).
DoorKey 6x6
DoorKey 8x8
Lava S9N1
Lava S9N2
Empty 16x16
AMRL
0.98 (0.005)
0.94 (0.005)
0.90 (0.009)
0.54 (0.412)
0.93 (0.009)
AMRLkard
0.84 (0.285)
0.02 (0.027)
0.86 (0.066)
0.02 (0.036)
0.93 (0.010)
KIAN
0.93 (0.005)
0.11 (0.293)
0.87 (0.021)
0.62 (0.237)
0.92 (0.007)
KoGuN
0.32 (0.359)
0.03 (0.040)
0.90 (0.011)
0.52 (0.441)
0.93 (0.009)
PPO
0.13 (0.200)
0.01 (0.018)
0.62 (0.357)
0.07 (0.156)
0.93 (0.011)
Table 3: Mean reward (one standard deviation) with different knowledge of varying informativeness levels. Agents on DoorKey 8x8 are trained for 300k frames and on Empty 16x16 for 200k frames. The details of knowledge modules are given in Appendix D AMRL consistently benefits from more informative modules, particularly achieving near-optimal performance in both environments. In contrast, AMRLhard, KIAN, and KoGuN exhibit limited improvements with increasing access to informative knowledge, suggesting that they struggle to fully leverage available knowledge.
DoorKey 8x8 medium
Empty 16x16 medium
low
high
low
high
AMRL
(0.03)
0.12 (0.28)
(0.01)
0.86 (0.23)
0.84 (0.22)
0.88 (0.18)
AMRLgard
0.02 (0.02)
0.03 (0.03)
0.07 (0.06)
0.65 (0.3)
0.71 (0.31)
0.65 (0.31)
KIAN
0.03 (0.04)
0.04 (0.06)
0.27 (0.39)
0.61 (0.35)
0.58 (0.38)
0.63 (0.34)
KoGuN
0.02 (0.03)
0.03 (0.028)
0.16 (0.31)
0.73 (0.3)
0.79 (0.3)
0.79 (0.27)
6.5 Module Informativeness
We investigate the effect of module informativeness, i.e, how informative it is for the task at hand as discussed in Subsection 4.1 further by training agents accessing modules with high, medium, and low informativeness (details in Appendix F. The results are reported in Table 3 We find that while agent with low module informativeness is able to solve the DoorKey 8x8 enviromment after 300k frames, AMRL with soft selection shows a slight advantage for medium informativeness and again significantly outperforms for high informativeness. On the Empty 16x16 environment low, medium and high knowledge has a positive impact on the agent's efficiency, but again AMRL bonefits most. The larger standard deviations are due to the fact that we have stopped training earlier (200k frames) than in Table 2 The training plots for these results can be found in Appendix
6.6 Robustness against Noisy Modules
Having verified the expected effects of differences in module informativeness for fixed number of modules, we now test the robustness of the selector against noisy modules. Specifically, we vary the density of the informative knowledge contained in all modules by adding in one, three, and five random modules, i.e., modules outputting equal action preferences for all actions. In Table 4 we report the achieved performance in terms of reward and Figure 5 shows the training runs on the LavaCrossing S9N1 environment. The training runs for the other environments can be found Appendix F.4
AMRL generally performs consistently well across environments, showing minimal degradation even as the number of random modules increases. AMRLiard shows poor performance under random modules in most environments. KIAN displays moderate robustness. Its performance degrades consistently as randomness increases. There is significant. variability in performance as indicated by the larger standard deviations in some environments (e.g., Door Key 8x8). KoGuN exhibits robustness similar to AMRL, showing consistently strong performance across most environments. Additionally, Figure shows that in contrast to AMRL which performs consistently, KoGuN performs well under moderate randomness, comparable to AMRL, but its sen- sitivity to higher levels of randomness (5 modules) reveals that it is more affected by high levels of randomness.
12

---

Published in Transactions on Machine Learning Research (11/2025)
Table 4: Evaluating robustness in presence of random modules. Agents on LavaCrossing S9N1 for 1.5 million frames, on Empty 16x16 for 150k frames, DoorKey 8x8 for 300k frames, DoorKey 6x6 for 200k frames. We report final reward =2 standard deviations across 10 random training seeds. AMRL is the most robust, closely followed by KoGuN. KIAN shows moderate robustness, and AMRL with hard selection is the most sensitive to random modules.
Lava S9N1 3
Empty 16x16 3
#random
0
5
5
AMRL
0.90
0.89
(0.014)
0.90 (0.011)
0.90 (0.012)
0.92
0.78
(0.264)
0.73 (0.290)
0.90 (0.077)
AMRLsard
0.86
0.86 (0.067)
0.21 (0.392)
0.01 (0.013)
0.60
0.22 (0.267)
0.14 (0.108)
0.06 (0.086)
KIAN
0.87
0.85 (0.026)
0.79 (0.064)
0.75 (0.057)
0.41
0.31 (0.234)
0.25 (0.077)
0.25 (0.104)
KoGuN
0.90
0.90 (0.009)
0.90 (0.016)
0.89 (0.025)
0.50
0.63 (0.341)
0.46 (0.303)
0.45 (0.281)
DoorKey 8x8
DoorKey 6x6
#random
3
5
3
5
AMRL
0.94
0.95 (0.003)
0.57 (0.510)
0.70 (0.404)
0.94
0.94
(0.004)
0.94 (0.004)
0.94 (0.006)
AMRLkard
0.04
0.03 (0.038)
0.01 (0.014)
0.02 (0.026)
0.94
0.75 (0.417)
0.44 (0.459)
0.02 (0.030)
KIAN
0.61
0.05 (0.069)
0.26 (0.444)
0.26 (0.376)
0.93
0.93 (0.004)
0.92 (0.006)
0.92 (0.005)
KoGuN
0.22
0.07 (0.040)
0.04 (0.032)
0.02 (0.024)
0.93
0.93 (0.005)
0.93 (0.005)
0.93 (0.006)
Agent: AMRL
Agent: AMRL hard
Agent: KIAN
Agent: KoGuN
66
Reward %
Rowardd
Roward
Reward
0
Frames
Frames
Frames
Frames
Figure 5: The reward dynamics during training on MiniGridd-Lvvarossinng frames). The set of original modules is modified by adding 1, 3, and 5 additional modules outputting uniformly random actions. AMRL demonstrates consistent robustness across varying numbers of random modules, while AMRLMards KTAN, and KoGuN exhibit higher sensitivity, with performance and sample efficiency degrading significantly as the number of random modules increases.
(training for 1.5 x 10t
Overall, AMRL is the most robust and resilient method when training in stochastic environments with random modules. It learns consistently and reaches near-optimal rewards regardless of the number of random modules. KoGuN has slightly worse performance but exhibits more variability, especially when the number of random modules is high. KIAN is moderately robust but less reliable, and AMRI-bard is confirmed to be sensitive randomness, almost entirely failing when random modules are introduced. This sensitivity of AMRLhard random modules and its failure to scale is direct consequence of the "all-or-nothing nature of hard selection. When the number of modules increases, especially with the introduction of uninformative, random ones, the selector's task becomes disproportionately more difficult. It is forced to make a single correct choice from a much larger and noisier set of options with an increasingly sparse signal.
Beyond the demonstration of robustness, the results presented in Table 4 and Figure 5 can be used for assessing the scaling propertics of the architecture as the number of modules increases from four modules to nine. AMRL scales generally well, while other methods struggle identify and utilize the relevant modules.
6.7 AMRL in Continuous State/Action Spaces
The evaluation on the OpenAL-Robotics- Gym shown in Figure 6 shows the success rate the FetchPickAndPPacce-t environment with continuous state and action spaces. AMRL achieves higher perfor- mance than both SAC and KIAN, reaching near-optimal success rates within the first 200k training steps and maintaining stability throughout. In contrast, SAC plateaus around 0.7 and KIAN around 0.5, indicating limited policy improvement. The results demonstrate that when paired with SAC, AMRL significantly en- hances learning efficiency and convergence, confirming its effectiveness in complex continuous control tasks.
13

---

Published in Transactions on Machine Learning Research (11/2025)
6.8 Limitations and Potential Extensions
DAC IAA
AMRL suffers from the following limitations. Firstly, it is worth noting that the complexity of the selector's action space grows with the number of modules; however, in practical settings, we expect that the number of modules will be limited. Secondly, naturally, when using existing knowledge its relevance and quality is crucial and AMRL is not an exception. However, we emphasize that diverse knowledge can contribute to developing more robust solutions, particularly in scenarios involving modules trained with noisy data or based on inconsistent rules.
3.0
Trannigg Stogs
206
Figure 6: Success rate on the PickAnd- Place environment.
In addition, AMRL may not always offer significant advantages over simpler modular RL approaches. For example, in the limit case in which all available modules are homogeneous RL policies trained under the same paradigm, the benefits of explicit heterogeneity are absent, and the additional overhead of managing a selector outweighs the gains. Similarly, in domains where high-quality heterogeneous knowledge is available, standard modular RL framework or even single- policy approach may be more practical. Future work could address these limitations by exploring adaptive mechanisms that decide when heterogeneity provides value, and by invest.gating scalability in scenarios with larger numbers of modules.
The release of new Large Language Models (LLM) have inspired a large body of work leveraging LLMs within RL agents, among others, provide intrinsic rewards (Klissarov al., 2023), guide exploration (Du et al, 2023 perform planning (Tchter et al. 2022 Huang ct al., 2022), and satisfy safety constraints (Yang et al. 2021 LLMs are indeed rich knowledge ledge base of human preferences and as such are compatible with our proposed approach. Future work could investigate incorporating modules based on LLMs, for example, to provide high level planning abilities and additional interpretability of the system. In fact, AMRL is designed as an open, flexible architecture, capable of supporting future learning mechanisms.
7 Conclusion
In this paper, we have rntuueed Augmented Modular Reinforcement Learning (AMRL), a framework that extends modular RL to explicitly support the integration of diverse heterogeneous knowledge, such as rules, trajectory datasets, and skills. AMRL handles heterogeneity at multiple levels: knowledge sources (rules, skills, retrieval, dynamic RL), state representations, and processing mechanisms. Through unified module abstraction and representation-agneet selector, it allows diverse modules to be composed seamlessly.
We have formally characterizod heterogencons knowledge and investigated several sources of information heterogeneity. We have implemented AMRL in environments from the Minigrid and the OpenAI-Robotie Gym suites, benchmarking against KTAN, KoGuN, PPO and SAC. Our results show that AMRL consistently improves sample efficiency, achieves safer exploration, and enhances robustness in the presence of noisy modules. Morcover, by structuring decision-making around modular and interpretable components, AMRL increases the transparency of agent behavior and facilitates generalization across tasks. We also analyzed hard versus soft selection, demonstrating that soft selection is particularly effective when integrating heterogencous modules.
Broader Impact Statement
AMRL is motivated by the goal of improving safety and robustness in reinforcement learning by enabling agents to incorporate heterogeneous knowledge such as rules, skills, and retrieval mechanisms. This allows for safer exploration, better interpretability, and more reliable behavior in domains where purely trial-and-crror learning can be unsafe. However, the safely and reliability of AMRL ultinately depend on the quality of the incorporated knowledge. Flawed, outdated, or biased rules and demonstrations could propagate harmful behaviors, and users may overestimate the guarantees provided by modular integration. While modular- ity enhances transparency relative to end-to-end RL, responsible deployment requires careful validation of knowledge sources and explicit consideration of ethical risks associated to it.
14

---

Published in Transactions on Machine Learning Research (11/2025)
Acknowledgments
Lorenz Wolf was supported by the UK Engincering and Physical Sciences Research Council (EP SSP/S02566/1) Mirco Musolesi acknowledges the support of the UK Engineering and Physical Sciences Research Council through the grant EP EP/X028569/L.
References
Pieter Abbeel, Adam Coates, and Andrew Y. Ng. Autonomous Helicopter Aerobatics through Apprenticeship Learning. International Journal of Robolics Research, 29(13):1608 11111-688-9,,,,,, 2010.
Jacob Andreas, Dan Klein, and Sergey Levine. Modular multitask reinforcement. learning with policy sketches. In Proceedings of the 34th International Conference on Machine Learning (ICML'17), 2017.
Andrew G. Barto and Sridhar Mahadevan. Recent Advancos in Hicrarchical Reinforcement Learning. Discrete Event Dynamic Systems, 13(4):341 379, 2003.
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding the mixture- of-experts layer in deep learning. In Advances in Neural Informalion Processing Systems (NeurIPS'2) 22)) 2022.
Maxime Chevalier-Boisvert, Lucas Willems, and Suman Pal. Minimalistic Gridworld Environment for Gym- nasium, 2018. URL https://githut Fttm///araameFuuuaaaaooo/a//aaa
Zih-Yun Chiu, Yi-Lin Tuan, William Yang Wang, and Michael C. Yip. Knowledge-Gruunee Reinforcement Learning, 2022. arXiv:2210.03729.
Zih-Yun Chiu, Yi-Lin Tuan, William Yang Wang, and Michael C. Yip. Flexible attention-based multi-policy fusion for efficient deep reinforcement learning. In Advances in Neural Information Processing Systems (N2arIPS'23), 2023.
Peter Dayan and GeofTrey E. Hinton. Feudal Reinforcement Learning. In Advances in Neural Information Processing Systems (NeurIPS'92), 1992.
Rodrigo de Lazcano, Kallinteris Andreas, Jun Jet Tai, Seungjae Ryan Leo, and Jordan Terry. Gymnasium robotics, 2024. URL http://github.
Coline Devin, Abhishek Gupta, Trevor Darrell, Pieter Abbeel, and Sergey Levine. Learning Modular Neural Network Policies for Multi-Tast and Multi-Robot Transfer. In Proceedings of the 2017 IEEE International Conference on Robotics and Automation (ICRA '17) 2017.
Bruce Digney. Learning hierarchical control structures for multiple tasks and changing environments. In Proceedings of the 5th Conference on the Simulation of Adaptive Behavior (5AB'98), 1998.
Yuqing Du, Olivia Watkins, Zihan Wang, Cédric Colas, Trevor Darrell, Pieter Abbeel, Abhishek Gupta, and Jacob Andreas. Guiding pretraining in reinforcement learning with large language models. In Proceedings of the 40th International Conference on Machine Learning (ICML 223) 2023.
Javier García and Fernando Fernández. A Comprehensive Survey on Safe Reinforcement Learning. Journal of Machine Learning Research, 16(42):1437-1480, 2015.
Dibya Ghosh, Jad Rahme, Aviral Kumar, Amy Zhang, Ryan P. Adams, and Sergey Levine. Why General- ization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability, In Advonces in Neural Information Processing Systems (NeurIPS'21 21), 2021.
Shivam Goel, Yash Shukla, Vasanth Sarathy, Matthias Scheutz, and Jivko Sinapov. Rapid-learm: A frame- work for learning to recover for handling novelties in open- wordd environments. In Proceedings of the 2022 IEFE International Conference on Development and Learning (TC)L,22), 2022.
15

---

Published in Transactions on Machine Learning Research (11/2025)
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bern- hard Schölkopf. Recurrent Independent Mechanisms. In Proceedings of the 9th International Conference on Representation Learning (ICLR 20), 2020.
Anirudh Goyal, Abram Friesen, Andrea Banino, Theophane Weber, Nan Rosemary Ke, Adrià Puigdomènech Badia, Arthur Guez, Mchdi Mirza, Peter C. Humphreys, Ksenia Konyushova, Michal Valko, Simon Osin- dero, Timothy Lillicrap, Nicolas Heess, and Charles Blundell. Netrieval-Augmenue Reinforcement Learn- ing. In Proceedings of the 39th International Conference on Machine Learning (ICML 222), June 2022.
Vaibhav Gupta; Daksh Anand, Praveen Paruchuri, and Akshat Kumar. Action selection for composable modular deep reinforcement learning. In Proceedings of the 20th International Conference on Autonomous Agents and MultiAgent Systems (AAMAS'2)), 2021.
Tuomas Haarnoja, Aurick Zhou, Picter Abbcol, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference on Machine Learning (TCML'18): 2018.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse control tasks through world models. Nature, 640(85559)6447-633) 2025.
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual Markov Decision Processes, 2015. arXiv:1502.02259,
Hosein Hasanbeig, Daniel Kroening, and Alessandro Abate. Certified reinforcement learning with logic guidance. Artificial Intelligence, 322:103949, 2023.
Bernhard Hengst. Hierarchical Reinforcement Learning. In Claude Sammut and Geoffrey Webb (eds.), Encyclopedia of Machine Learning, pp. 492-502. Springer US, Boston, MA, 2010.
Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. AdaRL: What, Where, and How to Adapt in Transfer Reinforcement Learning. In Proceedings of the 10th International Conference on Learning Representations (ICLR 22), 2022.
Peter Conway Humphreys, Arthur Guez, Olivier Tieleman, Laurent Sifre, Theophane Weber, and Timothy P. Lillicrap. Large-Scale Retrieval for Reinforcement Learning. In Advances in Neural Information Processing Systems (NeurIPS'22) 2022.
Brian Ichter, Anthony Brohan, Yevgen Chebotar, Chelsea Finn, Karol Hausman, Alexander Herzog, Daniel Ho, Julian Ibarz, Alex Irpan, Eric Jang, Ryan Julian, Dmitry Kalashnikov, Sergey Levine, Yao Lu, Car- olina Parada, Kanishka Rao, Pierre Sermanet, Alexander T. Toshev, Vincent Vanhoucke, Fei Xia, Ted Xiao, Peng Xu, Mengyuan Yan, Noah Brown, Michael Ahn, Omar Cortes, Nicolas Sievers, Clayton Tan, Sichun Xu, Diego Reyes, Jarek Rettinghouse, Jornell Quiambao, Peter Pastor, Linda Luu, Kuang- Huei Tee, Yuheng Kuang, Sally Jesmonth, Kyle Jeffrey, Rosario Jauregui Ruano, Jasmine Hsu, Keerthana Gopalakrishnan, Byron David, Andy Zeng, and Chuyuan Kelly Fu. Do as can, not as say: Ground- ing language in robotic affordances. In Proceedings of the 6th Annual Conference on Robol Learning (CoRL (22), 2022.
Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive Mixtures of Local Experts. Neural Computation, 3:79-87, 1991.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In Proceedings of the 5th International Conference on Learning Representations (ICLR'17) 2017.
Jeff Johnson, Matthijs Douze, and Hervé Jégon. Billion-scale similarity search with GPUs. IEEE Transac- tions on Big Data, 7(3), 2021.
B. Ravi Kiran, Tbrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A. Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. IEEE Transactions on Intelligent. Transportation Systems, (3(6):4909-4922) 2022.
16

---

Published in Transactions on Machine Learning Research (11/2025)
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A Survey of Zero-shot Generalisation in Deep Reinforcement Learning. Journal of Artificial Intelligence Research, 76:201-2000 264, 2023.
Martin Klissarov, Pierluca D'Oro, Shagun Sodhani, Roberta Raileanu, Pierre-Luc Bacon, Pascal Vincent, Amy Zhang, and Mikael Henaff. Motif: Intrinsic motivation from artificial intelligence feedback. In Proceedings of the 1st Workshop on Foundation Models for Decision Making at NeurIPS'23 2023.
Jens Kober, J Andrew Bagnell, and Jan Peters. Reinforcement Learning in Robotics: A Survey. International
Journal of Robotics Research, 32(11):1238 1274, September 2013.
Amarildo Likmeta, Alberto Maria Metelli, Andrea Tirinzoni, Riccardo Giol, Marcello Restelli, and Danilo Ro- mano. Combining reinforcement learning with rule-based controllers for transparent and general decision- making in autonomous driving. Robotics and Autonomous Systems, 131:103568, 2020.
Hsin- Yu Liu, Bharathan Balaji, Rajesh Gupta, and Dezhi Hong. Rule-basod policy regularization for rein- forcement learning-baser building control. In Proceedings of the 14th ACM International Conference Future Energy Systems (e-Energy 123), 2023.
Pierrick Lorang, Helmut Horvath, Tobias Kietreiber, Patrik Zips, Clemens Heitzinger, and Matthias Scheutz. Adapting to the "open world": The utility of hybrid hierarchical reinforcement learning and symbolic plan- ning. In Proceedings of the 2024 IEEE International Conference on Robotics and Automation (ICRA '24), 2024.
Jorge A. Mendez, Harm van Seijen, and Eric Eaton. Modular Lifelong Reinforcement Learning via Neural Composition. In Proceedings of the 10th International Conference Learning Representations (ICLR 222) 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7)4)):529-55 533, 2015.
Keerthiram Murugesan, Mattia Atzeni, Pavan Kapanipathi, Pushkar Shukla, Sadhana Kumaravel, Gerald Tesauro, Kartik Talamadupula, Mrinmaya Sachan, and Murray Campbell. Text-based RT. Agents with Commonsense Knowledge: New Challenges, Environments and Baselines. In Proceedings of the 35th AAAI Conference on Artificial Intelligence (AAAF21), 5 2021.
Ronald Parr and Stuart Russell. Reinforeement Learning with Hicrarchies of Machines. In Advances in Newral Information Processing Systems (NIPS'97), 1997.
Christian F. Perez, Felipe Petroski Such, and Theofanis Karaletsos. Generalized Hidden Parameter MDPs Transferable Model-based RL in Handful of Trials. In Proceedings of the 34th AAAI Conference on Artificial Inteligence (A)AI22)), 2020.
Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, Vikash Kumar, and Wojeiech Zaremba. Multi- Goal Reinforcement Learning: Challenging Robotics Environments and Request for Research, 2018. arXiv:1802.09464
Antonin Raffin. RL Basclines3 Zoo.
200 2020.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Basllne- Reliable reinforcement learning implementations. Journal of Machine Learning Re- search, 22(268):1-88 8, 2021.
Stuart Russell and Andrew L, Zimdars. Q-Decomposition for Reinforcement Learning Agents. In Proceedings of the 20th International Conference Machine Learning (ICML'03) '03), 2003.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimiza- tion algorithms, 2017. arXiv:1707.063:
17

---

Published in Transactions on Machine Learning Research (11/2025)
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Gooffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In Proceedings of the 5th International Conference on Learning Representations (ICLR'IP 17), 2017.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglon, Veda Panncorshclvam, Mare Lanctot, Sander Dicleman, Do- minik Grewe, John Nham, Nal Kalchbrenner, Ilya Sulskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7687)4484-88, Jamuary 2016.
Christopher Simpkins and Charles Isbell. Composable Modular Reinforcement Learning. In Proceedings of the 33rd AAAI Conference on Artificial Intelligence (AAAI'19), July 2019.
Nathan Sprague and Dana Ballard. Multiplo-Goal Reinforcement Learning with Modular SOrsa(O). In Proceedings of the 18th International Joinl Conference on Artificial Intelligence (AAAI'03), February 2003.
Richard S Sutton and Andrew G Barto. Reinforcement Learning: An Introduction. MIT Press, 2018
Richard S Sutton, Joseph Modayil, Michael Delp, Thomas Degris, Patrick M Pilarski, Adam While, and Doina Precup. Horde: A Scalable Real-time Architecture for Learning Knowledge from Unsupervised Sensorimotor Interaction. In Proceedings of the 10th International Conference on Autonomous Agents and Multiagent Systems (AAMAS'11), 2011.
Alexander Sasha Vezhnevets, Simon Osindero, Tom Schaul, Nicolas Heess, Max Jaderberg, David Silver, and Koray Kavukcuoglu. Feudal networks for hierarchical reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning (ICML' '17,, 2017.
Oriol Vinyals, Igor Babuschkin, Wojeiech M. Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H. Choi, Richard Powell, Timo Ewalds, Petko Georgiev, Junhyuk Oh, Dan Horgan, Manuel Kroiss, Ivo Danihelka, Aja Huang, Laurent Sifre, Trevor Cai, John P. Agapiou, Max Jaderberg, Alexander S. Vezhnevets, Rémi Leblond, Tobias Pohlen, Valentin Dalibard, David Budden, Yury Sulsky, James Molloy, Tom I.. Paine, Caglar Gulcchre, Ziyu Wang, Tobias Pfaff, Yuhuai Wu, Roman Ring, Dani Yogatama, Dario Wünsch, Katrina McKinney, Oliver Smith, Tom Schaul, Timothy Lillicrap, Koray Kavukcuoglu, Demis Hassabis, Chris Apps, and David Silver. Grandmaster level in StarCraft II using multi-agent reinforcement learning. Nature, 575(7782):335) (502))33333300 354, 2019.
Bernd Waschneck, André Reichstaller, Lenz Belzner, Thomas Altenmüller, Thomas Bauernhansl, Alexander Knapp, and Andreas Kyek. Deep reinforcement learning for semiconductor production scheduling. In Proceedings of the 29th Annual SEMI Advanced Semiconductor Manufacturing Conference (ASMC'18), pp. 301-306, 2018
Tsung-Yer Yen Yang, Michael Hu, Yinlam Chow, Peter Ramadge, and Karthik Narasimhan. Safe Reinforce- ment Learning with Natural Language Constraints, 2021. aIXiv:2010.05150.
Peng Zhang, Jianye Hao, Weixun Wang, Hongyao Tang, Yi Ma, Yihai Duan, and Yan Zheng. KoGuN: Accelerating Deep Reinforcement Learning via Integrating Human Suboptimal Knowledge. In Proceedings of the 29th International Joint Conference on Artificial Intelligence (IJCAI (0), 2020.
Victor Zhong, Tim Rocktäschel, and Edward Grefenstette. RTFM: Generalising to Novel Environment Dynamics via Reading. In Proceedings of the 8th International Conference Learning Regresentations (ICLR 20), 2020.
18

---

Published in Transactions on Machine Learning Research (11/2025)
A Policy Gradients for AMRL
A.1 Notation
Given a policy such that ((a s) = Pr(als) is the probability of taking action in state s when following policy and assuning we have an finite set of actions A [4 we use ((s) denote the vector with ith component being [s, Additionally, given list of modules M with corresponding policies {TM } MEM use (als) to denote the column vector with ith entry (als) where M is the ith module in M.
Furthermore, denote the selectors policy parameterized by as (-1s)) and the policy of the dynamic module dyn (-1s). Let the set M contain all modules of which one is the dynamic module.
A.2 Policy Gradients
The vanilla Policy Gradient of policy x is given by:
Σ Vo log n5 (ae 86) A= (se, ax)
VeJ (T²)
E Toxe
(6)
where J is the un-discounted finite time horizon reward to-go. Furthermore, the update rule used in PPO (Schulman et al. 2017) is given by:
0.+1 arg D TED,
ΣΣ
min
Tpp (az 8t) A (st, g (c, Ao (st, (s))
(7)
.TB (ar st)
While the below analysis is focused on the vanilla policy gradient, it justifies the use of PPO to train the selector.
A.3 Soft Selection
Using the specified notation we can write the AMRL policy with soft selection as:
7 (als) ग selector (s) (M
(8)
where M is the of modules, of which the only module parameterized is the dynamic module Magn with policy = dyn Alternatively, we can write:
π 0.0 (a)s) = (M ( (ST Meay (als) +
Σ
(9)
MEMiMaarr
where we have written out the inner product and split out the term corresponding to the dynamic module.
We can compute the policy gradients with respect to each of the parameter sets and 0 as follows. The gradient w.r.t. the dynamic module parameters 0 can be computed as:
VoJ(
E Toutt
Vo log x 0.0 (at st) An (s),
(10)
(11)
We can then derive:
Vo log 000 (at 8x)
Von (a+ 8x)
(12)
T4.
(a: st)
1
(Myn (ss) Vo (aels))
(13)
(a,
st)
19

---

Published in Transactions on Machine Learning Research (11/2025)
Similarly, compute the policy gradient w.r.t. the selector's parameters o, as follows:
V log n 0.0 (as $t) T0.0 (ac st)
V Vo
(at sc)
(14)
V u selector (s) [M(a)ss
(15)
T0.0 (ar st)
(v ( (8)) (aels)) +
Σ MEM
V (oT(AMss)TM(lll))
(16)
T0.9 (ar st)
(17)
A.4 Hard Selection
In the case of hard selection use the Gumbel Softmax trick to perform differentiable selection. In particular, want to sample from categorical distribution where the probabilities are given by the selector as (s). In other words, we select module M with probability equal to π T selector (M])s) We then obtain a sample vector y of dimensions |ssteactrl by setting:
exp (1log ( (Mils) + 91) /T)
Y; Σj=1 exp ((1og seleector (Mi M2]s)) + /T)
for i 1, I....,k,
(18)
where 01 ... 99 are i.dd samples drawn from Gumbel(O, 1). which can be casily sampled using an inverse transform. Note that the softmax function used as a continuous differentiable approximation. As the temperature t tends to 0, the samples tend to the samples of the true categorical distribution (aagg et al. 20177)
B Environments
In all the below environments the action space consists of 7 discrete actions shown in Table 5 observation size is 7x7 and is marked by the light grey grid centered around the agent Figure 7
The
Table 5: Action space of the Minigrid environments.
Num
Name
Action
0
left
Turn left
right
Turn right
2
forward
Move forward
3
pickup
Unused
4
drop
Unused
5
toggle
Unused
done Unused
C
Additional Implementation Details
For all agents below the PPO hyperpparametersssreee are set the default values provided in the rl-starter-file repository (attps://github.coll)) http://hhhh..cooooeeaaaaaaaaaaas
PPO. We use the network architecture of actor and critic network as implemented in the r1- rl-starter-ffles repository. In particular:
Embedding: A CNN with 3 convolutional layers produces an embedding of the image observations. Both the actor and the critic rely these image embeddings.
20

---

Published in Transactions on Machine Learning Research (11/2025)
i
(a) DoorKey 8x8
(b) LavaCrossing S9N2
Figure 7: Snapshots taken from two of the Minigrid environments.
Actor: The actor network is a layer fully connected neural network with layer width 64.
Critic:
The critic network is a 2 layer fully connected neural network with layer width 64.
AMRL. Tn the following, summarize the key design and implementation choices of algorithm:
Same as in PPO. The Embedding network is shared by the selector, the dynamic RL module, and the critic network.
Embedding:
The selector network is 2 layer fully connected neural network with layer width 64. It outputs a weight for each module. The softmax function is applied to the weights to map them to probabilities. To sample a module in the hard selection mechanism we use the Gumbel softmax distribution.
Selector:
Crilic: Same as in PPO.
KIAN. We rely on the original implementation provided by the authors: https://githuu..oomPaaaaaa KGRL
Embedding: A CNN with 3 convolutional layers produces an embedding of the image observations. The inner actor, the key network, and the critic rely on these image embeddings.
Actor: The actor consists of an internal actor, the query network, and key network. All 3 share an initial fully connected base layer of width 64, which is followed by an additional fully connected layer mapping from 64 features to the desired dimension (embedding size of 8 for the query and key network, and [A| for the internal actor). The keys modules are learned by the nn. Embedddin layer.
Crilic: Same as in PPO.
KoGuN, The following summarizes the key design choices of the KoGuN implementation.
Embedding: Same as for PPO. Both the actor, and the critic rely on these image embeddings.
The actor network is 2 layer fully connected neural network with layer width 64. It takes as imputs the output of the embedding network concatenated with the action preferences averaged across modules.
Actor:
Crilic: Same in PPO.
21

---

Published in Transactions on Machine Learning Research (11/2025)
Modules, All modules are implemented as described in the body of the main paper.
Compute Resources. The experiments were run on a CPU. No large amount of memory is required.
Computational Cost
An analysis of the computational cost of AMRL and baselines can be found in Table
Table 6: Analysis of the computational cost. The number of trainable parameters inchudes the trainable dynamic module for AMRL and the inner actor for KTAN. All knowledge based architeclures are comparable and rely the skill, retrieval and rule modules. The retrieval module searches for 4 neighbors in dataset of 3040 transitions. The faster inference times of PPO are caused by not calling the knowledge modules.
Method
# Trainable params
Total inference time (s)
PPO
19384
3.4 x 10-5
AMRL
19644
1.5 x 10
KoGuN
19832
1.5 x 10-1
KIAN
20448
1.6 x 10
E
Additional Experimental Details
DoorKey 8x8. The knowledge configurations of different levels used for the DoorKey 8x8 environment is as follows:
Low: Skill Empty 5x5, Retrieval Empty 5x5, Rules No rules.
Medium: Skill DoorKey 5x5, Retrieval DoorKey 5x5, Rules All rules.
High: Skill - DoorKey 6x6, Retrieval DoorKey 6x6, Rules - All rules.
Empty 16x16. On Empty, the knowledge for different informativeness levels is:
Low: Skill Empty 5x5, Retrieval Empty 5x5, Rules No rules.
Medium: Skill Empty 8x8, Retrieval Empty 8x8, Rules All rules.
High: Skill Empty 16x16, Retricval Empty 16x16, Rules All rules.
F Additional Results
F.1
Reward Plots
For completeness, present the training runs for the module informativeness experiment in Figures S and
F.2
Analyzing Learned Policies
To further investigate the learned policies and showcase improved interpretability, we plot the weights of the Selector in AMRL and the weights of the attention mechanism in KIAN to compare the learned policies in a range of selected environments. To obtain the weights, the agents are evaluated in the specified environments, weights are recorded for one episode (same configuration) from start to finish for all training seeds and we report mean weights together with =2 standard deviation confidence intervals. Note that weights are not comparable across different start positions and environment instances which is why we reported weights for the same episode. The extracted weights shown in Figure 10
22

---

Published in Transactions on Machine Learning Research (11/2025)
Agent: AMRL
Agent AMRL hard
Agent: KIAN
Agent: KoGuN
Reward
Rewward
Roward
LGw
High
-2
Frammes
Frames
Framess
Fromes
Figure 8: Average training run performance DoorKey 8x8 with different levels of knowledge.
Agent:
Agent AMRL
Agent:
Agent: KoGuN
Reward
Rewardd
Reward
Low
13
Medum
High
3000
50000
10000
Frames
Frames
Figure Average training run performance on Empty 16x16 with different of levels of knowledge, After 500k all agents have converged to roughly equal performance. We shorten the training runs and instead reported the performance after 200k frames of training.
F.3 Gumbel Temperature Ablation
We perform an ablation on the lemperature parameter of the Gumbel-Softmas distribution used for the hard selection mechanism. The results of the ablation study are reported for the following two cases: 1) the samples are discretized and the smooth approximation is used for gradionts as in the experiments reported in the paper; 2) the samples are not discretized.
For both cases we train for 100k frames on the DoorKey 6x6 environment and plot the mean rewards with (a) and 11 (b)
confidence intervals across 10 random training seeds. The results shown in Figures 11 for the discretized and not discretized cases respectively. We find that larger around 3 performs best in both cases. Additionally, the performance drops т decreases to 0 and the samples from the categorical distribution become more and more discrete regardless of whether we discretize samples to a one-hot vector in the forward pass.
F.4 Presence of Random Modules
we report the training curves on the DoorKey 8x8, DoorKey 6x6, and Empty 16x16 environments, for evaluating robustness in presence of random modules.
In Figures 12 and 14
23

---

Published in Transactions on Machine Learning Research (11/2025)
Dynamic: RL
Skil: Unlock
Retreval: Empty Random
Logic: All Rules
AMRL DoorKey_
KIAN DoorKey_ 6
0.5
Weight module
06
ghight of
04
0.3
02
0.0
5.0
7.5
10.0
12.5
15.0
17.5
15
20
AMRL DoorKey_ _
KIAN DoorKey _
12
Weight module
Weight of module
0.4
0.3
02
0.1
-02
Frames
Frames
AMRL LavaS9N1
KIAN LavaS9N1
1.0
0.8
Weight of module
0.
Woight of
02
-02
0.0
10
15
10
Frames
Frames
Figure 10: Weightings of modules in the final policy. Note that in contrast to AMRL, KIAN has not converged the DoorKey 8 8 environment and seems to be stuck because it puts too much weight on the retrieval module instead of the more relevant skill module. In the LavaCrossing S9N1 enviromment, AMRL has converged to almost fully rely on the dynamic RL module in contrast to KIAN, which has slightly more balanced weights.
24

---

Published in Transactions on Machine Learning Research (11/2025)
Discretization
No discretization
1.0
0.1.1
0.5
T.1.1
08
r 0.2
T= 0.2
Reward
7 0.5
Reward
== 0.5
c4
T=1
T=
62
T
T=6
0.0
T=
-02
-0.2
20000
40000
60000
50000
100000
20000
40000
80000
80000
100000
Frames
Frames
(a)
(b)
Figure 11: (a) Ablation on the temperature with discretization in the forward pass. (b) Ablation on the temperature without discretization in the forward pass.
Agent: AMRL
Agent AMRL hard
Agent:
Agent: KoGuN
0.
Reward
Roward
Rowarrd
Rooardd
5000 100000
10000 2000 20000 30000
1000
10000
WODs 2008 2030
1000
1000 100000 HSOH 2000
250000 Fraires
X000
15000 20000 20000 2000 Frames
Fremes
Fremes
Figure 12: The achieved reward logged throughout training Door Key 8x8. The set of original modules is modified by adding 1, 3, and 5 modules outputting uniformly random actions. It shows that AMRL can solve the environment with random modules present. However, training slower for 3 and random modules added. The results for other agents show significantly worse performance when random modules are present, but results are less conclusive due to significantly worse learning in the first place.
Agent: AMRL
Agent: AMRL hard
Agont: KIAN
Agent: KoGuN
1.2
1.0
Rewward
42
Roward
Rowward
25000 1508 1000 0000 R000M000000000
5000 7500 BOC.OOOORSOAAAAAAA
80000 7000 COROOOOCC 200000171300000000
2000 WOCRSORROOOTSSSIEEE
Fremes
Frames
Framess
Frames
Figure 13: The achieved reward logged throughout training on Door Key 6x6. is modified by adding 1, 3, and 5 modules outputting uniformly random actions.
The set of original modules
Agent: AMRL
Agent: AMRL_ hard
Agent: KIAN
Agent: KoGuN
Reward 02
6.
Roward
Reward
02
2000 #30 M000 S0I0 (E.8)
4.00 I0EU BUG0 TOUR 14823
200 BID 1A0000000111.
200C 4300 MIKO BOU W0000000011000
Fremes
Figure 14: The achieved reward logged throughout training on ShniGridd-Eppty166e2 training for 150k frames. The set of original modules modified by adding 1, 3, and modules outputting uniformly random actions.
25


Reviews:
[{"review_id": "NpNFQYsbln", "reviewer": "Reviewer_H5ND", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n1. Method looks interesting\n2. Paper easy to follow\n3. Sufficient ablation study\n\nCons:\n1. The experiments are for discrete action and thus finite state space, a experiment with continuous action/state space could provide a much stronger empirical support for the model.\n2. In Fig. 5, the AMRL hard looks worse in some cases, where 5 random modules result in flat reward, this should be further analyzed.\n3. I could be nitpicking, but the title looks a bit exaggerating. The proposed model expect each source of the \"heterogeneous knowledge\" to yield a corresponding policy. In practice, this may not be implemented easily, because human knowledge are not always \"well defined\".", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper propose a model to incorporate heterogeneous knowledge into the  Augmented Modular Reinforcement Learning method by establish policy $\\pi_M(s)$ for $M^{th}$ knowledge source and introduce a selector network that learns to decide the weight of each knowledge source based action that linearly combines to the final action.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Provide explanation for Fig. 5 AMRL failure cases.\n2. If possible, test some continuous state/action space cases.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "QBk0enmUna", "reviewer": "Reviewer_gHgt", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n* Clear motivation: The paper highlights the limitations of standard modular RL frameworks, which generally assume homogeneous RL-based sub-policies, and argues for the need to handle heterogeneous knowledge sources.\n* Unified framework: The proposed AMRL architecture provides a simple but flexible way to incorporate diverse module types (rules, retrieval, skills, dynamic RL) through a common selector. This framing emphasizes plug-and-play extensibility, which may inspire broader adoption.\n* Selector formalization: The paper systematically contrasts hard selection (Gumbel-softmax arbitration) and soft selection (weighted fusion), situating them within prior work on command arbitration and command fusion.\n* Empirical demonstration: Experiments on MiniGrid illustrate benefits in terms of sample efficiency, safety, and robustness, especially when combining informative heterogeneous modules. The use of safety metrics (e.g., unsafe actions in LavaCrossing) and interpretability analysis (selector weights) goes beyond simple reward reporting.\n* Discussion of extensions: The paper anticipates future directions, including the integration of LLM-based modules, which shows awareness of broader trends in RL + knowledge integration.\n\n### Weaknesses\n\n* Unclear novelty relative to prior HRL/modular RL:\n  The paper frames heterogeneous knowledge integration as something fundamentally new, but it is not convincingly argued why existing HRL or modular RL frameworks could not already accommodate sub-policies with different internal implementations. Traditional frameworks did not forbid heterogeneity; they simply assumed homogeneous policies for convenience. Thus, the conceptual novelty appears incremental.\n\n* Vague causal link between problem and solution:\n  The manuscript does not provide a precise explanation of what prevents prior methods from directly combining heterogeneous modules. The claimed “gap” is described mostly at a high level, without identifying specific theoretical or representational incompatibilities. The proposed solution — treating each module as producing an action preference and aggregating them via a selector — is reasonable but does not clearly resolve a previously insurmountable difficulty.\n\n* Ambiguity in state representation:\n  All modules ultimately consume the same state observation (MiniGrid grid view) and output distributions over the same action space. The heterogeneity lies only in their internal processing (rules, retrieval, RL policy), not in their input or output interfaces. This raises the question of whether AMRL introduces genuine heterogeneity or simply rephrases existing modular RL as a mixture of differently implemented policies. The paper does not explain how states would need to be transformed into different sub-states for modules that conceptually require distinct input formats.\n\n* Selector design is not fundamentally new:\n  The distinction between hard and soft selection largely mirrors well-studied paradigms (command arbitration vs. fusion, mixture-of-experts). The technical contribution of adopting Gumbel-softmax for differentiability and weighted averaging for fusion is modest, and their relationship to existing approaches should be acknowledged more explicitly.\n\n* Writing and positioning issues:\n  While the manuscript is well-polished (likely with LLM assistance), the narrative sometimes emphasizes breadth at the expense of causal clarity. The problem statement, contribution claims, and experimental design are not always tightly connected, leaving the reader uncertain whether AMRL solves a genuine limitation or simply demonstrates a repackaging of familiar ideas.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces Augmented Modular Reinforcement Learning (AMRL), a framework that extends traditional modular RL by integrating heterogeneous knowledge modules. Unlike standard approaches where sub-policies are typically homogeneous RL policies, AMRL allows modules to be derived from diverse sources, including rule-based logic, trajectory retrieval, pre-trained RL skills, and dynamic RL policies. A selector mechanism is proposed to arbitrate among these modules, either through hard selection (Gumbel-softmax sampling) or soft selection (weighted fusion).\n\nThe contributions claimed are:\n\n1. Formalization of heterogeneous modules for decision-making, providing a unified representation that accommodates different types of knowledge sources and processing mechanisms.\n2. Design of AMRL, an augmented modular RL framework with a selector that can flexibly combine heterogeneous modules in a plug-and-play manner.\n3. Empirical evaluation on MiniGrid environments, showing that AMRL with soft selection improves sample efficiency, safety, and robustness compared to baselines (PPO, KoGuN, KIAN).\n4. Discussion of how heterogeneity enhances interpretability and safety, and how the framework can potentially incorporate future knowledge sources such as LLMs.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "### Critical\n\n1. Clarify the causal motivation.\n   The current manuscript does not provide a precise explanation of why prior HRL/modular RL methods cannot already integrate heterogeneous sub-policies. The paper should articulate clearly:\n\n   * What practical or theoretical barriers exist in prior frameworks (e.g., comparability of reward scales, incompatibility of knowledge representations)?\n   * Why these barriers motivate AMRL specifically, rather than simply applying existing arbitration/fusion techniques.\n     Strengthening this causal chain is crucial to making the contribution well-justified.\n\n2. Clarify the role of state representations.\n   All modules in AMRL appear to consume the same environment state (e.g., 7×7 grid observation) and produce distributions over the same action space. If this is indeed the case, the heterogeneity is only in internal processing, not in state or action representation. The authors should explicitly discuss:\n\n   * How each module interprets the state, and whether any transformations or abstractions are applied.\n   * Why this setup should be regarded as genuinely heterogeneous rather than a collection of differently implemented policies.\n\n### Important but not Critical\n\n3. Position the novelty more carefully.\n   While the paper emphasizes heterogeneity as a new angle, prior HRL frameworks did not forbid heterogeneous modules; they simply did not emphasize them. The paper would be stronger if it acknowledged this, and reframed the novelty as:\n\n   * Providing a unified formalization of heterogeneous modules;\n   * Demonstrating empirically that such integration improves sample efficiency, safety, and robustness.\n     This more modest but accurate positioning would make the contribution clearer and more defensible.\n\n4. Refine discussion of hard vs. soft selection.\n   The current exposition could better situate these mechanisms in relation to existing work (e.g., command arbitration/fusion, mixture-of-experts). Rather than presenting them as new, the paper should highlight how they behave differently in the presence of heterogeneous modules, and why soft selection may be preferable.\n\n### Optional (would strengthen the paper)\n\n5. Expand discussion of limitations.\n   The paper already notes scaling issues with module count and reliance on knowledge quality. It would further help readers if the authors provided concrete scenarios where AMRL may not offer advantages over simpler modular RL (e.g., when all modules are homogeneous policies).\n\n6. Tighten writing for causal clarity.\n   Some passages emphasize breadth and polish over clarity (likely with automated editing). The introduction and related work would benefit from explicitly linking:\n\n   * Identified limitations of prior methods →\n   * Why those limitations matter in practice →\n   * How AMRL specifically addresses them.", "broader_impact_concerns": "no", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "n9JXsavgg4", "reviewer": "Reviewer_3asA", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths ###\n\n- The paper is well-written and straightforward to understand.\n- The proposed algorithm is simple and technically sound.\n- Experiments show that compared to baselines, the proposed method leads to higher reward and increased safety. It is also more robust to noisy sources of knowledge.\n\n### Weaknesses ###\n\n- Experiments were done on a fairly simple environment. GridWorld has a limited state space and action space.\n- It is not clear how rules could be incorporated in a more complex environment. For example, for a legged robot a rule could be to jump when there is an obstacle in front; how can this be translated into a policy recommendation?\n- No code is provided.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a modular / hierarchical reinforcement learning algorithm that is able to combine heterogeneous sources of information, such as rules, expert trajectories, and skills trained previously (on possibly unknown tasks). The main idea is to train a policy to select a module (or a mixture thereof) and then take the actions recommended by that module. The method is illustrated on different tasks in the GridWorld environment, showing improved reward learning compared to previous methods, increased safety, and interpretability.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "### Critical ###\n\n- Add experiments on a different environment, if possible.", "broader_impact_concerns": "No broader impact statement is provided. However, since the paper states that the increased safety due to the usage of rules is one of the main advantages of the algorithm, I feel that one should be included.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)