# Augmented Modular Reinforcement Learning based on Heterogeneous Knowledge

Lorenz Wolf
University College London
<EMAIL_ADDRESS>
&Mirco Musolesi
University College London
University of Bologna
<EMAIL_ADDRESS>

###### Abstract

In order to mitigate some of the inefficiencies of Reinforcement Learning
(RL), modular approaches composing different decision-making policies to
derive agents capable of performing a variety of tasks have been proposed. The
modules at the basis of these architectures are generally reusable, also
allowing for "plug-and-play" integration. However, such solutions still lack
the ability to process and integrate multiple types of information
(knowledge), such as rules, sub-goals, and skills. We propose Augmented
Modular Reinforcement Learning (AMRL) to address these limitations. This new
framework uses an arbitrator to select heterogeneous modules and seamlessly
incorporate different types of knowledge. Additionally, we introduce a
variation of the selection mechanism, namely the Memory-Augmented Arbitrator,
which adds the capability of exploiting temporal information. We evaluate the
proposed mechanisms on established as well as new environments and benchmark
them against prominent deep RL algorithms. Our results demonstrate the
performance improvements that can be achieved by augmenting traditional
modular RL with other forms of heterogeneous knowledge.

##  1 Introduction

Reinforcement learning (RL) has been a cornerstone in the field of machine
learning, with its capability of optimizing decision-making processes in
complex dynamic environments [50, 31]. This has been demonstrated in a myriad
of applications, such as robotics [32], dynamic resource allocation [10, 55],
and complex game playing [38, 46, 54]. Yet, despite these advances, the
applicability to real world problems and practical deployment of RL solutions
are hindered by issues concerning sample efficiency and safety concerns [15].
In fact, current model-free RL algorithms may require hundreds of million
samples to converge to the optimal policy and exploration may cause the agent
to take risky actions which in safety-critical situations could lead to
catastrophic outcomes.

Modular RL aims to mitigate sample inefficiency by combining several modules
implementing different "skills" solving sub-problems, such that together they
solve more complex tasks [44, 49, 47]. In general, modular approaches provide
flexibility and can speed up learning when skills previously acquired can be
transferred to solve new tasks [11]. However, such solutions still lack the
ability to process and integrate multiple types of information, such as rules,
sub-goals, and skills. We will refer to the latter with the term knowledge in
the remainder of this work.

In many cases, existing prior knowledge and other sources of information could
speed up the learning process and improve safety, however, current solutions
are not flexible and generalizable. Indeed, efforts have been made to
incorporate different sources of information into RL, most notably in form of
text manuals and video game-play [58, 14] but also knowledge graphs [40].
However, these are mostly designed to address a specific problem and limited
to one source of information, lacking at the same time the seamless
integration of diverse types and the flexibility to acquire new knowledge.
Prior knowledge has also been used to mitigate safety concerns, for example to
reduce the number of catastrophic events during learning and exploration [1,
48, 16, 15]. While successfully reducing unsafe actions in specific
environments, these methods suffer from the same set of limitations, since
they mostly consider knowledge in form of demonstrations or teacher advice
[15].

Figure 1: AMRL architecture: the framework is able to access several sources
of heterogeneous knowledge via modules. The modules can then be updated based
on the environment feedback.

To address these limitations, we propose Augmented Modular Reinforcement
Learning (AMRL), an approach that flexibly incorporates diverse sources of
knowledge, enhancing the learning process and improving safety in RL systems
(Figure 1). Building upon the framework of modular RL, we augment it through
the seamless integration of various types of knowledge. Our approach not only
includes modular skills but also incorporates rules, logic, and other forms of
structured and/or declarative knowledge that can guide the agent’s behavior
more effectively. This augmentation not only strengthens the modularity of the
RL system but also facilitates faster learning and ensures safer operations.
In our design, a general selector (or arbitrator) will have the role of
selecting among standard RL modules and others that are based on other forms
of heterogeneous knowledge. The selector itself is agnostic with respect to
the actual knowledge representation. This allows for greater adaptability in
diverse scenarios, making it suitable for a variety of applications. In this
work we make the following main contributions:

  * •

Formalization of a theoretical framework for the integration of heterogeneous
modules for decision making. We introduce a formalism for designing
heterogeneous modules for decision making, enabling the integration of diverse
types of knowledge representations. This provides a unified framework for
expressing knowledge as rules, logic, skills (implemented as RL modules), and
so on, which can be flexibly combined to tackle complex tasks.

  * •

Augmentation of modular RL to incorporate heterogeneous knowledge. We discuss
the design of AMRL as an augmentation of a classic modular RL framework that
enables the seamless integration of heterogeneous knowledge. By incorporating
it into the modular RL framework, we enhance the learning process, allowing
the agent to leverage prior knowledge for faster convergence, but also acquire
and update knowledge for adaptability.

  * •

Extension of the selection mechanism in command arbitration. We propose an
extension to the selection mechanism. The Memory-Augmented Arbitrator, brings
progress made on memory and temporal abstraction in RL to modular RL and, more
generally, architectures with a selection mechanism.

  * •

Evaluation of the proposed architecture and extensions. We evaluate the
proposed architecture on a challenging custom environment, characterized by
partial observability and multi-modality, as well as a set of existing
environments involving a variety of design aspects, including adaptability,
efficiency, and safety.

##  2 Related work

##### Hierarchical RL.

Hierarchical RL (HRL) decomposes a problem into several smaller subtasks
organized in a hierarchy[12]. 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[24]. Each
task is learned by a policy and at decision time, the higher level policy
triggers a policy in the level below. The goal is for the agent to perform
task abstraction by focusing on more high-level tasks, so that it will suffer
less from the curse of dimensionality [5]. This is a very active area of
research with some of the earlier works dating back to the 1990s [9, 53, 42].
One of the most interesting and popular proposals is that based on options
[52].

##### Modular RL.

HRL is closely related to modular RL, which aims to decompose the learning
process into complementary policies / modules that can be combined by a
controller [51, 18, 28]. 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
[21, 44, 36, 3]. In this work we primarily focus on command arbitration.111It
is worth noting that the approach presented in this paper could be extended to
command fusion. We plan to investigate this aspect in future work. Finally, it
is worth noting that in previous work on reward decomposition, it is assumed
that the overall reward function or the value function is the sum of all
individual reward functions or value functions ([35] and [45] respectively).
Our approach is more general, since it is not based on that assumption.

##### Command arbitration.

In the command arbitration paradigm an arbitrator decides which module takes
control over the agent in a given state; the selected module then executes an
action. Earlier works include GM-Sarsa [49] and Q-decomposition [44]. More
recent methods Arbi-Q [47] and GRACIAS [21] assume that the arbitrator
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[47], while the authors of GM-Sarsa and Q-Decomposition
originally reported performance reductions when utilizing off-policy
techniques due to future actions potentially being chosen sub-optimally[49].

##### Knowledge augmented reinforcement learning.

Grounding RL agents with external knowledge is not a new area of research per
se. For example, progress has been made on incorporating knowledge provided to
the agent via text for example as manual for a game [58] or knowledge graphs
connecting concepts and their characteristics [40]. These works are related
but different in that text features and state features are treated jointly to
produce one action rather than using different modules to propose an action
and then performing action selection. We leverage modularity for flexibility,
robustness, and interpretability. Recent advances in Natural Language
Processing (NLP) and the release of new Large Language Models (LLM) have
inspired a large body of work leveraging LLMs within RL agents, among others,
to guide exploration [13], perform planning [6, 26], satisfy safety
constraints [57], and interact [41]. Another promising line of research
focuses on equipping the agent with a retrieval mechanism to search over a
memory of trajectories in order to inform the current decision [19, 27]. In
both works updating the knowledge is not formally considered and the knowledge
is restricted to trajectories. Most similar in motivation to our work are
Knowledge Grounded Reinforcement Learning agents (KGRL), which use an
attention mechanism to combine a learning agent with pre-trained policies
$\pi_{g_{i}}:\mathcal{S}\rightarrow\mathcal{A}$ (knowledge) [8]. Importantly,
the knowledge cannot be updated online through experience and is only
composable w.r.t. the weighting assigned to the different policies. Instead,
we build on hierarchical modular RL, allowing us to share experience between
modules and update knowledge if needed. Consequently, our approach is more
general and flexible while still maintaining simplicity.

##  3 Preliminaries

Let
$\mathcal{M}=\left(\mathcal{S},\mathcal{A},\mathcal{P},R,\rho_{0},\gamma,T\right)$
represents a discrete-time finite-horizon discounted Markov decision process
(MDP). $\mathcal{M}$ consists of the following elements: the state space
$\mathcal{S}$, the action space $\mathcal{A}$, a transition probability
distribution
$\mathcal{P}:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow\mathbb{R}_{+}$,
a reward function $R:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$, an
initial state distribution $\rho_{0}:\mathcal{S}\rightarrow\mathbb{R}_{+}$, a
discount factor $\gamma\in[0,1]$, and the horizon $T$. Generally, the policy
$\pi_{\theta}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}_{+}$ is
parameterized by $\theta$ which we optimize in order to maximize the expected
discounted return under the policy. Formally, we maximize
$\eta\left(\pi_{\theta}\right)=\mathbb{E}_{\pi_{\theta}}\left[\sum_{t=0}^{T}\gamma^{t}R\left(s_{t},a_{t}\right)|s_{t},a_{t}\right]$,
where

$s_{0}\sim\rho_{0}\left(s_{0}\right),a_{t}\sim\pi_{\theta}\left(a_{t}\mid
s_{t}\right)\text{, and }s_{t+1}\sim\mathcal{P}\left(s_{t+1}\mid
s_{t},a_{t}\right).$ (1)

In this work we use Q-Learning to train an RL agent [56]. We define the
optimal value function
$Q^{*}(s,a)=\max_{\pi_{\theta}}\eta\left(\pi_{\theta}\right).$ At optimization
step $j$, $\theta_{j}$ is updated by minimizing the loss function

$L_{j}\left(\theta_{j}\right)=\mathbb{E}_{s,a\sim\rho(\cdot)}\left[\left(y_{j}-Q\left(s,a;\theta_{j}\right)\right)^{2}\right],$
(2)

where
$y_{j}=\mathbb{E}_{s^{\prime}\sim\mathcal{P}}\left[r+\gamma\max_{a^{\prime}}Q\left(s^{\prime},a^{\prime};\theta_{j-1}\right)\mid
s,a\right]$ and $\rho(\cdot)$ is the behavior distribution [39]. Finally, 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). A POMDP is represented by the tuple
$\left(\mathcal{S},\mathcal{A},\mathcal{P},\Omega,\mathcal{O},R,\rho_{0},\gamma,T\right)$,
where in addition to the already known quantities we have introduced the
observation space $\Omega$ the observation probability distribution
$\mathcal{O}(\cdot)$. Instead of observing $s\in\mathcal{S}$, the agent now
observes $o\in\Omega$ which is generated from the state $s$ via
$o\sim\mathcal{O}(s)$.

##  4 Augmented modular RL based on heterogeneous knowledge

In this section we outline our proposed architecture for the integration of
diverse types of knowledge. We first present how the decision-making process
of an agent can be augmented by means of knowledge; we then discuss potential
types of knowledge, which can be exploited by agents; finally, we present the
overall architecture based on an extension of hierarchical modular RL.

###  4.1 Heterogeneous knowledge

Given the fact that our framework is based on RL, the focus is on sequential
decision making. Hence we either require representations of knowledge that are
"actionable" (e.g., an action can be selected using it) or we need learn how
to make the knowledge actionable. Consequently, natural choices are procedural
or declarative knowledge, for example in form of a world model [22] or a
database of past trajectories [19, 27]. While these might be the most
effective representations for RL agents, it can be hard to obtain these forms
of knowledge since it requires the ability to deploy an agent to learn and
acquire them in the first place. For example, when learning a world model we
need to explore the entire state action space freely, and for trajectory
retrieval we either require an offline dataset or the ability to collect large
amounts of data without any knowledge incorporated.

It is therefore important to flexibly use and acquire knowledge that is more
abstract in nature, for instance skills (for example, in the form of a trained
model for a MDP), text, or other information available within the environment
or externally. We also note the fact that knowledge sources such as text,
logic or beliefs are difficult to exploit for agents but they are easy to
express, validate and understand for humans. As such, augmenting RL agents
with this type of heterogeneous knowledge also creates a strong opportunity to
make them safer [15] and more interpretable [17].

###  4.2 Architecture

The AMRL framework is based on modular hierarchical RL. In Modular RL the
modules represent policies receiving rewards from the environment [44] or
pseudo-rewards specified by the developer [51]. Here we extend the modules to
be heterogeneous and thus enable the agent to incorporate various types of
knowledge as described above. In AMRL the modularity not only provides
flexibility, but it is also crucial for facilitating integration within a
hierarchical RL framework. The proposed architecture consists of the following
main components: modules $M_{i}$ for $i=1,\cdots,N$ with $N$ the number of
modules, which form the foundation, and a selector, which combines the modules
together as shown in Figure 2. We will now describe both components in more
detail.

####  4.2.1 Modules

Intuitively, one can think of a module in the modular RL context as a skill, a
policy solving the subtask at hand. In contrast to previous works in modular
RL, the modules here are heterogeneous. They are not limited to RL policies
and they can be diverse in terms of how they form preferences over actions. A
module could encode action preferences based on the knowledge of logic or
rules such as "do not step forward if $X$" or "pick up the object $Y$ when it
is in front of you". More formally, a module $M_{i}$ implements a generic
mapping $m_{i}$ from some input space $\Phi_{i}$ to the action space
$m_{i}:\Phi_{i}\rightarrow A_{i}$, where $A_{i}\subseteq\mathcal{A}$. The
input space $\Phi_{i}$ can be static or dynamic.

Figure 2: High-level architecture of an AMRL agent consisting of modules
$M_{i}$ with internal mappings $m_{i}$ and the selector. Inputs to the modules
are $\phi_{i}\in\Phi_{i}$ which do not necessarily depend on time, and
feedback signals $r_{i,t}$ for the modules and $r_{\star,t}$ for the selector.
At time t the selector chooses module $M_{\star}$ and the action $a_{\star,t}$
is performed.

Given an input $\phi_{i}\in\Phi_{i}$ module $M_{i}$ proposes an action
$a_{i,t}\in A_{i}\;\forall i$.222The subscript in $\phi_{i}$ indicates the
input space and for simplicity of notation we write $a_{i,t}$ instead of
$a^{module}_{i,t}$. We note that the modules are not required to have the same
action spaces. However, this can restrict experience sharing between the
modules since if the performed action $a_{\star,t}$ is proposed by another
module and $a_{\star,t}\notin\mathcal{A}_{i}$ module $M_{i}$ is not able to
learn from the current transition. From now on we assume
$A_{i}=\mathcal{A}\>\forall i$.

Modules and knowledge can be fixed or dynamic. If the knowledge and,
consequently, the mapping of a module $m_{i}$ is to be updated during
training, we require the input space to be observable over time and grounded
in the environmentsuch that the effect of actions proposed by $M_{i}$ can be
observed in the environment and used to update the knowledge. This is not a
necessary condition as $\Phi_{i}$ can be static and exist outside the
environment, for example in form of an external knowledge base. When updating
the knowledge, or learning it, a feedback signal is required. This can be the
global reward $r_{\star,t}$ to obtain task specific knowledge or a measure of
the effectiveness of actions based on the knowledge.

Similarly to modular RL, we may specify modular rewards
$\\{r_{i,t}\\}_{i=1}^{N}$, where $r_{i,t}$ is the reward for module $M_{i}$ at
time $t$, to favor modular learning. If possible, specifying modular rewards
can help learn a given decomposition of the knowledge and the task, but it is
not required.

####  4.2.2 Selector

The selector is the high-level controller/arbitrator that combines the
modules, by choosing when to trigger which one, to perform a more complex
task. We would like to point out again that the selector itself is agnostic
with respect to the actual knowledge representation.

At each time step, it receives the current observation $s_{t}$ from the
environment and selects one of the modules. The proposed action of the
selected module is then executed by the agent. Thus, the selector can be
formalised as a policy $a^{selector}_{t}\sim\pi_{selector}(\cdot|s_{t})$ over
the discrete action space $\mathcal{A}_{selector}=\\{M_{1},\dots,M_{N}\\}$, or
equivalently as a value function $Q_{selector}(s_{t},a_{t},\theta_{selector})$
parameterized by $\theta_{selector}$. By the definition of
$\mathcal{A}_{selector}$ we have that

$M_{\star}=\underset{M_{k}}{\operatorname{argmax}}\;Q_{selector}(s_{t},M_{k};\theta_{selector})\quad\text{and}\quad
a_{\star,t}=m_{*}(\phi_{\star})$ (3)

In the vanilla architecture shown in Figure 2, the optimal value function of
the selector is approximated by a Deep-Q-Network (DQN) [39, 38] trained with
experience replay.333Note that the general framework is not limited to
Q-Learning and can also be used with on-policy methods. At time $t$ the
selector receives the global reward $r_{\star,t}$, which corresponds to
solving the main task. Introducing a global reward avoids potential issues
related to the adoption of varying reward scales and off-policy learning [47].
Similar to previous works in Modular RL, and in contrast to others on reward
decomposition, we do not require $r_{\star,t}=\sum_{i=1}^{N}r_{i,t}$. In fact,
all modules and the selector may receive the same reward signal.

To make a decision at time $t$, each module expresses a preference and
proposes an action $a_{i,t}$ based on $\phi_{i}$. The selector observes
$s_{t}$ and chooses one of the modules, which we refer to as $M_{\star}$, to
be triggered. Consequently, action $a_{\star,t}$ is performed. Then the
modules to be updated and the selector receive feedback $r_{i,t}$ and
$r_{\star,t}$, respectively, based on the state action pair
$\left(s_{t},a_{\star,t}\right)$ and are updated correspondingly. Note how
this allows us to share experience between modules, which means they can
benefit from each other and collectively speed up learning. We provide the
pseudo-code for this process in Algorithm 1.

Algorithm 1 Training of an Augmented-Modular RL agent.

$\\{M_{i}\\}_{i=1}^{N}$, $Q_{selector}(\cdot,\cdot;\theta_{selector})$

for each episode do

$s_{0}\leftarrow\text{Initial state}$

for each timestep $t$ do

$a_{i,t}\leftarrow M_{i}(\phi_{i})$ $\triangleright$ Get proposed preferences

$M_{\star}\leftarrow\underset{M_{k}}{\operatorname{argmax}}\;Q_{selector}(s_{t},M_{k};\theta_{selector})$
$\triangleright$ Select a module

$a_{\star,t}\leftarrow\left[a_{1,t},\cdots,a_{N,t}\right]$ $\triangleright$
Get action from $M_{\star}$ and execute

Observe $s_{t+1}$ and rewards $r_{\star,t},\\{r_{i,t}\\}_{i=1}^{N}$

Update modules $M_{i}$ based on $r_{i,t}$ and Selector based on $r_{\star,t}$

end for

end for

##  5 Memory-Augmented Arbitrator

In this section we discuss an architecture extension for command arbitration
methods in general. We equip the arbitrator with a memory and temporal
awareness. We refer to this solution as Memory-Augmented Arbitrator (Mem-
Arbi). We believe that this arbitrator, even if it is essentially only an
extension with respect to the core architecture presented in the previous
sections, can be seen as general reusable component and technique for
advancing the design of modular RL solutions. In the following, we will
discuss the design in detail.

(a) Vanilla architecture

(b) Memory-Augmented Arbitrator

Figure 3: Architecture comparison showing the main differences between the
vanilla AMRL in Figure 3(a) and the Memory-Augmented Arbitrator, which relies
on temporal information through its hidden state $h_{t}$, in Figure 3(b).

We propose to equip the arbitrator with a memory of past selections (Figure
3(b)). In other words, the arbitrator should be able to plan ahead and be
aware of previously chosen modules. This will not only help to make more
consistent selections but also mitigate potential limitations of the modules
and knowledge. In order to achieve temporal abstraction in the module
selection we equip the arbitrator with memory, specifically by using a Long
Short Term Memory (LSTM) recurrent neural network [25]. In the past, recurrent
neural networks have successfully been deployed in RL (e.g., [4, 37, 20, 29],
particularly to overcome partial-observability [23], but to the best of our
knowledge, not in the context of command arbitration. We argue that, in many
cases, it is sensible to take past module selections into account. Even in
cases when this is not possible, for instance due to modules being pre-trained
or the available knowledge not being suitable for temporal abstraction, Mem-
Arbi can exploit information from previous observations during the selection.
This is especially useful if the partial-observability can be overcome by
using the knowledge accessed by a module. Mem-Arbi also brings command
arbitration closer to the research on options[52].

##  6 Evaluation

The evaluation focuses on the agents ability to acquire new knowledge and
utilize prior knowledge in several different situations.

###  6.1 Environments

We evaluate our proposed architecture on new and existing environments, each
of which presents a specific set of challenges to the agents. We will use them
to show the advantages in adopting AMRL as framework for decision making. In
the following, we will describe the key characteristics of each environment. A
full description of these environments for reproducibility is available in the
Supplementary Material B of this article. We report screenshots of these
environment in Figure 4.

(a) Collect

(b) Lava-Crossing

(c) Door-Key

Figure 4: The environments used for evaluation.

##### Collect.

This custom environment is a collection task in a Gridworld (Figure 4(a)). The
agent’s task is to collect all balls placed in the environment. To increase
the difficulty we combine partial observability with strong multi-modality in
the sampling distribution of the balls. In our experiments the modes for each
episode are sampled uniformly on the whole grid. The user can specify a
minimum inter-mode distance to ensure a minimum level of multi-modality in
every episode (implemented via rejection sampling). We then sample locations
from a multi-modal Gaussian distribution with the modes derived in the
previous step and bucket the locations into squares on the grid. In our
experiments we specify a grid size of $15\times 15$, a view size of $5\times
5$, and $2$ modes with a distance of at least $10$ between them, unless
specified otherwise. 12 balls are randomly allocated to the 2 modes and placed
around them as described. Here the knowledge can be the locations of the
modes, as well as prior knowledge, such as rules like "pick up a ball when it
is in front of you" or knowing how to go to the modes.

##### Lava-Crossing.

To investigate potential benefits of the proposed approach with respect to
safety constraints, we consider an environment introduced in "Safety
Gridworlds" [34]. In particular, we train and evaluate our agent on `MiniGrid-
LavaCrossingS9N2-v0` [7]. Here the agent has to navigate through 2 crossings
in the lava stream to reach its goal. When the agent falls in the lava the
episode ends, i.e., ensuring safe exploration is fundamental in this
environment. Using modular rewards we can learn the fact that stepping on lava
is fatal, yielding a module to play it safe and avoid lava. This is also one
of the prior pieces of knowledge given to the agent. As an additional type of
prior knowledge, we also consider a module that simply explores new places.

##### Door-Key.

To investigate the agent’s ability to learn and incorporate specific skills we
use the `MiniGrid-DoorKey-8x8-v0` environment [7]. To reach the goal, the
agent has to first collect the key to then unlock a door which is blocking the
way. This environment is suitable to study whether the agent can acquire
skills such as collecting a key and using them to achieve the goal.

###  6.2 Implementation and training

All agents are implemented in PyTorch[43]. The hyper-parameters are tuned with
Bayesian Optimization provided by Optuna [2]. The agents perform
$\epsilon$-greedy exploration with exponential $\epsilon$-decay. The states
for all environments are encoded with a convolutional neural network (CNN)
[33]. The state representations produced by the CNN are then further processed
by fully connected layers. In the case of Mem-Arbi we use instead an LSTM
layer. The network parameters are optimized with the Adam Optimizer [30]. We
use a maximum number of $150$ steps per episode unless specified otherwise.
This makes learning much harder, since many episodes do not reach termination
and hence no reward, but it allows us to better understand the effect of
integrating and acquiring knowledge. The reported scores are a moving average
over the last $30$ episodes. The implemented baselines are a DQN agent [39,
38] and a DRQN one [23]. Additionally, in the Collect environment we also
evaluate against a DQN policy combining the state observation with the
information about the location of the modes. The full architecture details and
hyper-parameter settings can be found in the Supplementary Material C.

###  6.3 Acquiring new knowledge

In the first set of experiments we investigate the agent’s ability to learn
new knowledge and the benefits gained by doing so. In particular, we consider
the AMRL agent giving it simply the environment reward AMRL as well as the sum
of module rewards, which we denote by AMRL and AMRL-ds, respectively. In the
$15\times 15$ collect task AMRL and AMRL-ds perform similarly well, and learn
quicker than the baselines, closely followed by the DQN+k agent relying on the
position of the modes and the current observation jointly (Figure 5(a)).
Importantly, the DQN agent does not yield two separate modules, which can
later be reused for other tasks but just one. The purely state based DQN
baseline is less stable, and it is slower in terms of learning. The DRQN agent
reaches the best overall score but also takes the longest to learn due to the
complexity added by the LSTM layer. Note that both DQN+k as well as DRQN
achieve similar levels of performance indicating that the knowledge can be
used to overcome partial observability. However, none of the agents manage to
perfectly solve the environment, with 150 steps per episode.

(a) Collect

(b) Door-Key

(c) Lava-Crossing

Figure 5: Agents acquiring knowledge in the different environments. AMRL and
AMRL-ds differ only in the reward given to the selector, AMRL-ds receives the
sum of the module rewards.

In the Door-Key (Figure 5(b)) and the Lava-Crossing (Figure 5(c)) environments
we notice a clear difference in performance. While both AMRL and AMRL-ds
obtain rewards early on and clearly outperform both DQN and DRQN. The
baselines are struggling to reach the goal within 150 steps per episode and
consequently learn much slower or do not learn at all.

###  6.4 Using prior knowledge

Using prior knowledge is just as important as acquiring new knowledge. We
provide the AMRL agent with important pieces of prior knowledge. In the
collection task, we adopt a DQN trained to solve the task reasonably well
(collecting 10/12 balls on average), i.e., an oracle; in the lava crossing
environment we report the results for the declarative knowledge "avoid lava";
and for the Door-Key task we consider the subgoal "collect the key first". For
the collection task we have increased the size of the grid to $20\times 20$
with a minimum inter-mode distance of $15$ and allow the agent 200 steps per
episode to increase the difficulty. We note that to fully solve the
environments, longer training is required; however, shorter training already
reveals clear advantages of the AMRL agents. The results confirm what we
observe in the case of knowledge acquisition. In the larger Collect
environment we reach quicker convergence (Figure 6(a)) and in the Lava-
Crossing and Door-Key tasks AMRL agents learn significantly quicker than the
baselines (Figures 6(b), 6(c) and 6(d)).

###  6.5 Evaluating the Memory-Augmented Arbitrator

To test whether Mem-Arbi is indeed beneficial even when modules do not rely on
a temporal abstraction, we evaluated it with prior knowledge on all three
environments, although for Collect and Door-Key memory is more important. On
the Collect task AMRL+Mem, AMRL with Mem-Arbi, yields significant improvements
over the vanilla AMRL agent (Figure 6(a)). On the Door-Key task, we observe a
similar improvement in terms of learning performance (Figure 6(b)). However,
we do not observe convergence and AMRL+Mem appears to be less stable than
vanilla AMRL. On Lava-Crossing AMRL+Mem does not significantly enhance the
learning process (Figure 6(d) and 6(c)). We attribute this result to the fact
that memory plays a less important role in this environment.

(a) Collect with "Oracle" knowledge

(b) Door-Key with "Get Key" knowledge

(c) Crossing with "Explore" knowledge

(d) Crossing with "Avoid Lava" knowledge

Figure 6: Performance in Collect, Door-Key and Lava-Crossing with
heterogeneous knowledge.

###  6.6 Discussion and limitations

Our evaluation on the Collect, Door-Key, and Lava-Crossing environments shows
clear advantages of AMRL over the baselines. AMRL is able to acquire new
knowledge and outperform other agents thanks to it. Most noticeable are the
improvements in performance and convergence speed when using relevant prior
knowledge. On memory heavy tasks, the proposed extension AMRL+Mem improves on
AMRL confirming that the selection mechanism is able to compensate for modules
without temporal abstraction. The comparison between AMRL and AMRL-ds shows
that the latter may be more effective for acquiring knowledge, since it
targets all modules; however, weighting the modular rewards sensibly may
require domain knowledge and careful consideration. It is also worth noting
that the complexity of the arbitrator state space grows with the number of
modules; however, in practical settings, we expect that the number of modules
will be limited. Naturally, when using existing knowledge, the relevance and
quality of said knowledge is crucial. This is also the case for AMRL.
Consequently, AMRL with "Oracle" knowledge for Collect and "Get Key" knowledge
for Door-Key results in stronger improvements than AMRL on Lava-Crossing.
"Explore" knowledge causes the agent to step on lava more frequently while
"Avoid Lava" makes it overly cautious.

##  7 Conclusion

In this paper we have discussed the design, implementation and evaluation of
Augmented Modular Reinforcement Learning (AMRL), a solution that is able to
process and integrate multiple types of information (knowledge), such as
rules, sub-goals, and skills for more efficient and effective decision-making.
The proposed framework uses an arbitrator to select heterogeneous modules and
seamlessly incorporate different types of knowledge. This arbitrator is
agnostic with respect to the knowledge representation itself. We have also
introduced a variation of the selection mechanism, namely the Memory-Augmented
Arbitrator, which adds the capability of exploiting temporal information. We
have evaluated the proposed mechanisms on established as well as new
environments and benchmark them against prominent deep RL algorithms. Our
results show significant performance improvements that can be achieved by
augmenting traditional modular RL with heterogeneous knowledge.

## References

  * [1] Pieter Abbeel, Adam Coates and Andrew Y. Ng  “Autonomous Helicopter Aerobatics through Apprenticeship Learning”  In _International Journal of Robotics Research_ 29.13, 2010, pp. 1608–1639
  * [2] Takuya Akiba et al.  “Optuna: A Next-generation Hyperparameter Optimization Framework”  In _Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining (SIGKDD’19)_ Anchorage AK USA: ACM, 2019, pp. 2623–2631
  * [3] 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
  * [4] Bram Bakker  “Reinforcement Learning with Long Short-Term Memory”  In _Proceedings of Advances in Neural Information Processing Systems (NeurIPS’01)_ 14 MIT Press, 2001
  * [5] Andrew G. Barto and Sridhar Mahadevan  “Recent Advances in Hierarchical Reinforcement Learning”  In _Discrete Event Dynamic Systems_ 13.4, 2003, pp. 341–379
  * [6] Anthony Brohan et al.  “Do as I can, not as I say: Grounding language in robotic affordances”  In _Conference on Robot Learning 2023 (CORL’23)_ , 2023
  * [7] Maxime Chevalier-Boisvert, Lucas Willems and Suman Pal  “Minimalistic Gridworld Environment for Gymnasium”, 2018  URL: https://github.com/Farama-Foundation/Minigrid
  * [8] Zih-Yun Chiu, Yi-Lin Tuan, William Yang Wang and Michael C. Yip  “Knowledge-Grounded Reinforcement Learning” arXiv:2210.03729 [cs]  arXiv, 2022
  * [9] Peter Dayan and Geoffrey E Hinton  “Feudal Reinforcement Learning”  In _Advances in Neural Information Processing Systems (NeurIPS’92)_ , 1992
  * [10] “DeepMind AI Reduces Google Data Centre Cooling Bill by 40%”  URL: https://www.deepmind.com/blog/deepmind-ai-reduces-google-data-centre-cooling-bill-by-40
  * [11] Coline Devin et al.  “Learning Modular Neural Network Policies for Multi-Task and Multi-Robot Transfer”  In _Proceedings of the 2017 IEEE International Conference on Robotics and Automation (ICRA’17)_ Singapore: IEEE Press, 2017, pp. 2169–2176
  * [12] Bruce Digney  “Learning hierarchical control structures for multiple tasks and changing environments”  In _Proceedings of the 5th Conference on the Simulation of Adaptive Behavior (SAB’98)_ 98, 1998
  * [13] Yuqing Du et al.  “Guiding Pretraining in Reinforcement Learning with Large Language Models” arXiv:2302.06692 [cs], 2023
  * [14] Linxi Fan et al.  “MineDojo: Building Open-Ended Embodied Agents with Internet-Scale Knowledge”  In _Advances in Neural Information Processing Systems (NeurIPS’22) - Datasets and Benchmarks Track_ , 2022
  * [15] Javier García and Fernando Fernández  “A Comprehensive Survey on Safe Reinforcement Learning”  In _Journal of Machine Learning Research_ 16.42, 2015, pp. 1437–1480
  * [16] Javier García and Fernando Fernández  “Safe Exploration of State and Action Spaces in Reinforcement Learning”  In _Journal of Artificial Intelligence Research_ 45, 2012, pp. 515–564
  * [17] Claire Glanois et al.  “A Survey on Interpretable Reinforcement Learning” arXiv:2112.13112 [cs], 2022
  * [18] Anirudh Goyal et al.  “Recurrent Independent Mechanisms”  In _Proceedings of the 9th International Conference on Representation Learning (ICLR’20)_ , 2020
  * [19] Anirudh Goyal et al.  “Retrieval-Augmented Reinforcement Learning”  In _Proceedings of the 39th International Conference on Machine Learning (ICML’22)_ , 2022
  * [20] Audrunas Gruslys et al.  “The Reactor: A fast and sample-efficient Actor-Critic agent for Reinforcement Learning”  In _Proceedings of the 6th International Conference on Machine Learning (ICML’18)_ , 2022
  * [21] 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’21)_ , 2021
  * [22] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba and Timothy Lillicrap  “Mastering Diverse Domains through World Models” arXiv:2301.04104 [cs, stat], 2023
  * [23] Matthew Hausknecht and Peter Stone  “Deep Recurrent Q-Learning for Partially Observable MDPs”  In _Proceedings of the 2015 AAAI Fall Symposium_ , 2015
  * [24] Bernhard Hengst  “Hierarchical Reinforcement Learning”  In _Encyclopedia of Machine Learning_ Boston, MA: Springer US, 2010, pp. 495–502
  * [25] Sepp Hochreiter and Jürgen Schmidhuber  “Long Short-Term Memory”  In _Neural Computation_ 9.8, 1997, pp. 1735–1780
  * [26] Biwei Huang et al.  “AdaRL: What, Where, and How to Adapt in Transfer Reinforcement Learning”  In _Proceedings of the 10th International Conference on Learning Representations (ICLR’22)_ , 2022
  * [27] Peter Conway Humphreys et al.  “Large-Scale Retrieval for Reinforcement Learning”  In _Advances in Neural Information Processing Systems (NeurIPS’22)_ , 2022
  * [28] Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan and Geoffrey E. Hinton  “Adaptive Mixtures of Local Experts”  In _Neural Computation_ 3, 1991, pp. 79–87
  * [29] Steven Kapturowski et al.  “Recurrent Experience Replay in Distributed Reinforcement Learning”  In _Proceedings of the 7th International Conference on Learning Representations (ICLR’19)_ , 2019
  * [30] Diederik P. Kingma and Jimmy Ba  “Adam: A Method for Stochastic Optimization”  In _Proceedings of 3rd International Conference for Learning Representations (ICLR’15)_ , 2017
  * [31] B. Kiran et al.  “Deep Reinforcement Learning for Autonomous Driving: A Survey” arXiv:2002.00444 [cs], 2021
  * [32] Jens Kober, J Andrew Bagnell and Jan Peters  “Reinforcement Learning in Robotics: A Survey”  In _International Journal of Robotics Research_ 32.11, 2013, pp. 1238–1274
  * [33] Y. Lecun, L. Bottou, Y. Bengio and P. Haffner  “Gradient-based learning applied to document recognition”  In _Proceedings of the IEEE_ 86.11, 1998, pp. 2278–2324
  * [34] Jan Leike et al.  “AI Safety Gridworlds” arXiv:1711.09883 [cs], 2017
  * [35] Zichuan Lin et al.  “Distributional Reward Decomposition for Reinforcement Learning”  In _Advances in Neural Information Processing Systems (NeurIPS’19)_ , 2019
  * [36] Jorge A. Mendez, Harm Seijen and Eric Eaton  “Modular Lifelong Reinforcement Learning via Neural Composition”  In _Proceedings of the 10th International Conference on Learning Representations (ICLR’22)_ , 2022
  * [37] Volodymyr Mnih et al.  “Asynchronous Methods for Deep Reinforcement Learning”  In _Proceedings of the 33rd International Conference on Machine Learning (ICML’16)_ , 2016
  * [38] Volodymyr Mnih et al.  “Human-level control through deep reinforcement learning”  In _Nature_ 518.7540 Nature Publishing Group, 2015, pp. 529–533
  * [39] Volodymyr Mnih et al.  “Playing Atari with Deep Reinforcement Learning”  In _Proceedings of the NIPS Deep Learning Workshop 2013_ , 2013
  * [40] Keerthiram Murugesan et al.  “Text-based RL Agents with Commonsense Knowledge: New Challenges, Environments and Baselines”  In _Proceedings of the AAAI Conference on Artificial Intelligence (AAAI’21)_ , 2021
  * [41] Joon Sung Park et al.  “Generative Agents: Interactive Simulacra of Human Behavior” arXiv:2304.03442 [cs], 2023
  * [42] Ronald Parr and Stuart Russell  “Reinforcement Learning with Hierarchies of Machines”  In _Advances in Neural Information Processing Systems (NeurIPS’97)_ 10 MIT Press, 1997
  * [43] Adam Paszke et al.  “Automatic differentiation in PyTorch”  In _NIPS-W_ , 2017
  * [44] Stuart Russell and Andrew L Zimdars  “Q-Decomposition for Reinforcement Learning Agents”  In _Proceedings of the 20th International Conference on Machine Learning (ICML’03)_ , 2003
  * [45] Harm Seijen et al.  “Hybrid Reward Architecture for Reinforcement Learning”  In _Advances in Neural Information Processing Systems (NIPS’17)_ , 2017
  * [46] David Silver et al.  “Mastering the game of Go with deep neural networks and tree search”  In _Nature_ 529.7587 Nature Publishing Group, 2016, pp. 484–489
  * [47] Christopher Simpkins and Charles Isbell  “Composable Modular Reinforcement Learning”  In _Proceedings of the 33rd AAAI Conference on Artificial Intelligence (AAAI’19)_ , 2019
  * [48] Yong Song, Yi-bin Li, Cai-hong Li and Gui-fang Zhang  “An Efficient Initialization Approach of Q-learning for Mobile Robots”  In _International Journal of Control, Automation and Systems_ 10.1, 2012, pp. 166–172
  * [49] Nathan Sprague and Dana Ballard  “Multiple-Goal Reinforcement Learning with Modular Sarsa(0)”  In _Proceedings of the 18th International Joint Conference on Artificial Intelligence (AAAI’03)_ , 2003
  * [50] Richard S Sutton and Andrew G Barto  “Reinforcement Learning: An Introduction”  MIT Press, 2018
  * [51] Richard S Sutton et al.  “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
  * [52] Richard S. Sutton, Doina Precup and Satinder Singh  “Between MDPs and semi-MDPs: A framework for temporal abstraction in reinforcement learning”  In _Artificial Intelligence_ 112.1-2, 1999, pp. 181–211
  * [53] Alexander Sasha Vezhnevets et al.  “FeUdal Networks for Hierarchical Reinforcement Learning”  In _Proceedings of the 34th International Conference on Machine Learning (ICML’17)_ , 2017
  * [54] Oriol Vinyals et al.  “Grandmaster level in StarCraft II using multi-agent reinforcement learning”  In _Nature_ 575.7782, 2019, pp. 350–354
  * [55] Bernd Waschneck et al.  “Deep reinforcement learning for semiconductor production scheduling”  In _Proceedings of the 29th Annual SEMI Advanced Semiconductor Manufacturing Conference (ASMC’18)_ , 2018, pp. 301–306
  * [56] Christopher JCH Watkins and Peter Dayan  “Q-learning”  In _Machine Learning_ 8 Springer, 1992, pp. 279–292
  * [57] Tsung-Yen Yang et al.  “Safe Reinforcement Learning with Natural Language Constraints” arXiv:2010.05150 [cs], 2021
  * [58] Victor Zhong, Tim Rocktäschel and Edward Grefenstette  “RTFM: Generalising to Novel Environment Dynamics via Reading”  In _Proceedings of the 8th International Conference on Learning Representations (ICLR’20)_ , 2020

##  Appendix A Memory-Augmented Arbitrator

In this section we describe the architecture of the Memory-Augmented
Arbitrator as well as the training process in more detail. In particular, for
the latter we follow the "Bootstrapped Random Updates" strategy [23].

###  A.1 Experience sampling

During learning, we require batches containing sequences of at least min_seq
consecutive transitions in the environment. Instead of storing transitions we
store episodes with a length of at least min_seq in the replay buffer. To
sample training experience we sample batch_size episodes and randomly sample a
sequence of the specified length from each of those episodes. One of the
downsides of this strategy is that we are no longer sampling sequences
uniformly from the replay buffer, but rather the episodes are sampled
uniformly and then sequences are sampled uniformly within the episodes.
Consequently, a sequence in a shorter episode is more likely to be selected.
While there are downsides to this strategy, mainly zeroing the hidden state of
the LSTM layer at each update, it has been shown to perform competitively with
other strategies [23]. We note that it is important to strike a balance
between sequence length allowing the agent to learn from information further
in the past and the speed at which we fill up the replay buffer. The hidden
states are reset to zero at the beginning of each sequence.

###  A.2 Decision making

At the beginning of each episode the agent’s hidden state is set to zero.
Then, at decision time $t$, Mem-Arbi and DRQN take inputs $(s_{t},h_{t})$ and
output $(a^{selector}_{t},h_{t+1})$ or $(a_{t},h_{t+1})$, respectively. First,
the current state $s_{t}$ is encoded via a CNN and then passed to the LSTM
layer together with the hidden state $h_{t}$. The LSTM outputs the new hidden
state $h_{t+1}$ and the processed state, which is passed through fully
connected layers to obtain the estimated Q-values. The architecture only
differs from that shown in Figure 7 in that layers Linear-8 and Linear-9 are
replaced by an LSTM layer with specified sequence length and hidden state
size.

##  Appendix B Implementation details and hyperparameters of the environments

###  B.1 Collect

The action space for the Collect environment is given by
$\mathcal{A}=\\{\text{left},\text{right},\text{forward},\text{pickup}\\}$. An
instance of the environment is generated as described in the main article.
Furthermore, the reward for successfully picking up a ball is set to $1$.

###  B.2 Door-Key and Lava-Crossing

The action spaces and rewards for Door-Key and Lava-Crossing are as specified
in the environment documentation.444Mingrid documentation:
https://minigrid.farama.org/environments/minigrid. The action space for Lava-
Crossing is $\mathcal{A}=\\{\text{left},\text{right},\text{forward}\\}$ and
for Door-Key it is given by
$\mathcal{A}=\\{\text{left},\text{right},\text{forward},\text{pickup},\text{toggle}\\}$.

In both environments, a reward of

$1-0.9*\frac{step\\_count}{max\\_steps}$

is given for success, and $0$ for failure.

##  Appendix C Experimental settings for reproducibility

Figure 7: Example architecture of the DQN agent

###  C.1 Architectures

An example of of the DQN architecture for a batch size of 64 and an agent with
a view size of 7 is shown in Figure 7. The CNN block used for the Collect
environment is smaller since the agent has a more limited view size. The
selector is itself a network of this form and the width of the fully connected
layers is denoted by "Latent dim". Similarly, for the DQN and DRQN baselines
this hyperparameter is used to set the width of the hidden layers; for DQN+k
it sets the width of the hidden layers processing the state representation.

In the case of DQN+k, the baseline based on the training of a DQN policy over
mode locations and the state representation, each of those inputs is encoded
separately and the state and knowledge embeddings are then concatenated and
passed through further fully connected layers to produce the estimated
Q-values. Thus, we have the following additional hyperparameters:

  * •

Combined latent dimension: the width of the hidden fully connected layer after
concatenation;

  * •

Knowledge latent dimension: the width of the hidden fully connected layer
encoding the mode locations;

  * •

Knowledge embedding size: the size of the knowledge embedding (e.g. after
processing the mode locations);

  * •

State embedding size: the size of the state embedding.

The LSTM layer of the selector in AMRL+Mem and the DRQN baseline also have the
length of the sequences fed into the LSTM and the size of the hidden state of
the LSTM as additional hyperparameters.

###  C.2 Hyperparameter settings

For all experiments we specify the maximum number of steps per episode as 150,
except for Collect with "Oracle" knowledge, where we set the maximum number of
steps per episode to 200. For all agents and all environments we used a
discount factor of 0.99.

####  C.2.1 Learning knowledge

The selected hyperparameters for the learning knowledge experiment in the
Collect environment are reported in Table 1 and in Table 2 for the Lava-
Crossing and Door-Key environments.

Hyperparameter | DQN | DQN+k | DRQN | AMRL | AMRL-ds
---|---|---|---|---|---
Epsilon decay | $0.995$ | $0.98$ | $0.98$ | $0.982$ | $0.982$
Module learning rate | $0.0008$ | $0.0008$ | $0.0008$ | $0.0036$ | $0.0.0036$
Selector learning rate | - | - | - | $0.0076$ | $0.0076$
Latent dimension | $32$ | $64$ | $64$ | $16$ | $16$
Combined latent dimension | - | $16$ | - | - | -
Knowledge latent dimension | - | $64$ | - | - | -
Knowledge embedding size | - | $8$ | - | - | -
State embedding size | - | $40$ | - | - | -
Batch size | $64$ | $128$ | $128$ | $128$ | $128$
Buffer size | $10^{5}$ | $5\times 10^{4}$ | $5\times 10^{4}$ | $5\times 10^{4}$ | $5\times 10^{4}$
Update frequency | $8$ | $8$ | $8$ | $4$ | $4$
Soft update | $0.0023$ | $0.0023$ | $0023$ | $0.0012$ | $0.0012$
LSTM hidden size | - | - | $16$ | - | -
LSTM sequence length | - | - | $8$ | - | -

Table 1: Hyperparameter settings for learning knowledge in the Collect environment. Hyperparameter | DQN | DRQN | AMRL | AMRL-ds
---|---|---|---|---
Epsilon decay | $0.995$ | $0.98$ | $0.992$ | $0.992$
Module learning rate | $0.0008$ | $0.0008$ | $0.0036$ | $0.0036$
Selector learning rate | - | - | $0.0077$ | $0.0077$
Latent dimension | $32$ | $64$ | $64$ | $64$
Batch size | $64$ | $128$ | $128$ | $128$
Buffer size | $10^{5}$ | $10^{5}$ | $10^{5}$ | $10^{5}$
Update frequency | $8$ | $8$ | $4$ | $4$
Soft update | $10^{-3}$ | $0.0023$ | $0.0012$ | $0.0012$
LSTM hidden size | - | $10$ | - | -
LSTM sequence length | - | $16$ | - | -

Table 2: Hyperparameter settings for learning knowledge in the Lava-Crossing
and Door-Key environments. Note that, due to the difficulty of the task, the
baselines perform similarly for a variety of hyperparameter values.

####  C.2.2 Using prior knowledge

The selected hyperparameters for the experiment using "Oracle" prior knowledge
in the Collect environment are reported in Table 3, those for "Avoid Lava" and
"Explore" knowledge in the Lava-Crossing environment are presented in Tables 4
and 5, respectively. The hyperparameters for using "Get Key" knowledge in the
Door-Key environment are recorded in Table 6.

Hyperparameter | DQN | DQN+k | AMRL | AMRL+Mem
---|---|---|---|---
Epsilon decay | $0.996$ | $0.996$ | $0.996$ | $0.996$
Module learning rate | $0.0036$ | $0.003$ | $0.0036$ | $0.0057$
Selector learning rate | - | - | $0.0007$ | $0.00033$
Latent dimension | $64$ | $64$ | $64$ | $32$
Combined latent dimension | - | $16$ | - | -
Knowledge latent dimension | - | $64$ | - | -
Knowledge embedding size | - | $8$ | - | -
State embedding size | - | $40$ | - | -
Batch size | $128$ | $128$ | $64$ | $64$
Buffer size | $10^{5}$ | $10^{5}$ | $10^{5}$ | $10^{4}$
Update frequency | $4$ | $8$ | $4$ | $4$
Soft update | $0012$ | $0.0023$ | $0.0012$ | $0.0071$
LSTM hidden size | - | - | - | $64$
LSTM sequence length | - | - | - | $10$

Table 3: Hyperparameter settings for using "Oracle" prior knowledge in the Collect environment. Hyperparameter | DQN | DRQN | AMRL | AMRL+Mem
---|---|---|---|---
Epsilon decay | $0.99$ | $0.995$ | $0.992$ | $0.997$
Module learning rate | $0.0008$ | $0.0008$ | $0.0036$ | $0.00576$
Selector learning rate | - | - | $0.0077$ | $0.00053$
Latent dimension | $64$ | $64$ | $64$ | $64$
Batch size | $64$ | $64$ | $64$ | $64$
Buffer size | $10^{5}$ | $10^{5}$ | $10^{5}$ | $10^{4}$
Update frequency | $8$ | $8$ | $4$ | $8$
Soft update | $10^{-3}$ | $0.0023$ | $0.0012$ | $0.0071$
LSTM hidden size | - | $10$ | - | $16$
LSTM sequence length | - | $16$ | - | $8$

Table 4: Hyperparameter settings for "Avoid Lava" prior knowledge in the Lava-Crossing environment. Hyperparameter | DQN | DRQN | AMRL | AMRL+Mem
---|---|---|---|---
Epsilon decay | $0.99$ | $0.995$ | $0.992$ | $0.997$
Module learning rate | $0.0008$ | $0.0008$ | $0.0036$ | $0.00576$
Selector learning rate | - | - | $0.0077$ | $0.00053$
Latent dimension | $64$ | $64$ | $64$ | $64$
Batch size | $64$ | $64$ | $128$ | $64$
Buffer size | $10^{5}$ | $10^{5}$ | $10^{5}$ | $10^{4}$
Update frequency | $8$ | $8$ | $4$ | $8$
Soft update | $10^{-3}$ | $0.0023$ | $0.0012$ | $0.0071$
LSTM hidden size | - | $10$ | - | $16$
LSTM sequence length | - | $16$ | - | $8$

Table 5: Hyperparameter settings for "Explore" prior knowledge in the Lava-Crossing environment. Hyperparameter | DQN | DRQN | AMRL | AMRL+Mem
---|---|---|---|---
Epsilon decay | $0.999$ | $0.999$ | $0.999$ | $0.999$
Module learning rate | $0.0008$ | $0.0008$ | $0.00086$ | $0.00057$
Selector learning rate | - | - | $0.00016$ | $0.000132$
Latent dimension | $64$ | $64$ | $64$ | $32$
Batch size | $64$ | $128$ | $128$ | $82$
Buffer size | $10^{5}$ | $10^{5}$ | $10^{5}$ | $10^{4}$
Update frequency | $8$ | $8$ | $8$ | $4$
Soft update | $10^{-3}$ | $0.0023$ | $0.0012$ | $0.0071$
LSTM hidden size | - | $100$ | - | $200$
LSTM sequence length | - | $32$ | - | $15$

Table 6: Hyperparameter settings for "Get Key" prior knowledge in the Door-Key
environment.

###  C.3 Oracle training

For comprehensiveness, we also report the training dynamics of the oracle,
which is used for the Collect prior knowledge experiment, in Figure 8.

Figure 8: Training of the oracle on a $20\times 20$ Collect task (minimum
inter-mode distance of $15$) with a maximum of $200$ steps per episode.