Build generally reusable agent-environment interaction models
Abstract
This paper tackles the problem of how to pre-train a model and make it generally reusable backbones for downstream task learning.
In pre-training, we propose a method that builds an agent-environment interaction model by learning domain invariant successor features from the agent’s vast experiences covering various tasks, then discretize them into behavior prototypes which result in an embodied set structure.
To make the model generally reusable for downstream task learning, we propose (1) embodied feature projection that retains previous knowledge by projecting the new task’s observation-action pair to the embodied set structure and (2) projected Bellman updates which add learning plasticity for the new task setting.
We provide preliminary results that show downstream task learning based on a pre-trained embodied set structure can handle unseen changes in task objectives, environmental dynamics and sensor modalities
Introduction
d e a b c
Can we pre-train a generally reusable backbone to accelerate downstream task reinforcement learning (RL)?
While such backbones are popular in computer vision and natural language processing (e.g., pre-trained ResNet (He et al., 2016) and pretrained GPT3 (Brown et al., 2020), pre-trained decision-making backbones that support general downstream task learning are a few.
The challenges largely attribute to the vast complexity and variety of decision-making tasks, that the pre-trained model will suffer from the out-of-distribution nightmare when reused to learn a new task.
For example, a downstream task’s environment dynamics or task objectives may be unseen during pre-training.
Or we may need to add an extra sensor modality (e.g., language) to the agent to best fit the task requirement (e.g., understanding human commands).
How can we pre-train a backbone that handles the above-unseen changes in downstream task learning, without making any assumptions on the task similarity (Duan et al., 2016; Gupta et al., 2018) of environment similarity (Packer et al., 2018; Rajeswaran et al., 2016)?
Empirically, humans are capable of reusing experience from one task to another, and the tasks we can handle are surprisingly heterogeneous (Fig. 1), from playing blocks, hockey and karate, to learning a carpenter or cooking job.
It’s natural to ask, (1) what model is generally shared between heterogeneous tasks?
And (2) how can it fit in any downstream task learning with unseen settings?
Systematically, thinking on the above questions results in two perspectives of viewing what component is shared across tasks—the world and the agent (Harutyunyan, 2020).
The first one seeks to build a Bayesian world model (Ha & Schmidhuber, 2018), with which, the agent stores prior knowledge about the physical world.
For example, “balls will roll” and “walking on ice is slippery”.
Then the Bayesian agent selects a proper world configuration for downstream task learning.
However, since downstream tasks can be complex and diverse, then how large and how accurate the world model should be is not clear.
The second perspective seeks to build an embedded agency (Orseau & Ring, 2012), that the agent gradually constructs knowledge learned from its past experience, respecting the fact that the world is far more complex than the agent.
The constructed knowledge can be representations or models and will accelerate downstream task learning.
The embedded agency perspective can also gradually build a world model for the agent, but unlike the Bayesian world model , the agent’s past experience shapes what the world model is, how large it is, and how accurate it is.
This paper relates to the second viewpoint.
Though the embedded agency view is conceptually appealing, technically how to build it and the promising benefits w.r.t.
sample efficiency or continual learning capability are not clear.
Without being too ambitious to build a complete embedded agency, in this paper, we focus on a more practical starter: how the agent constructs a model from its embodied experience, and how it accelerates new task learning when environments, task objectives, and even sensor modalities are unseen .
Specifically, we format the problem of building a generally reusable model as an embodied set construction approach that constructs a domain invariant set structure using a large dataset covering the agent’s vast task experiences.
We call it an agent-environment interaction model since the set elements are discretized successor features (Barreto et al., 2017) which represent how a decisionmaking agent interacts with the world.
We call our approach embodied set construction since the set structure is constructed so that it is a shared component across tasks and environments, which can be viewed as an embodiment of the agent.
We call it generally reusable since learning by projecting onto the embodied set structure enables learning stability-plasticity.
Our contributions are as follows:
• We propose an embodied set construction approach to pre-train a reusable agent-environment interaction model by firstly learning domain invariant successor features from the agent’s past experience, then discretizing them to behaviour prototypes which result in an embodied set structure.
Note that the constructed embodied set structure, though optimized to be domain-invariant, still suffers the out-of-distribution (O.O.D.) problem when reused in downstream tasks under unseen settings.
• To make the pre-trained model generally reusable, we propose two projection-based techniques that enable learning stability-plasticity: (1) embodied feature projection which retains previous knowledge by projecting the new task’s observation-action pairs onto the embodied set structure, and (2) projected Bellman update which adds plasticity for the new task learning.
• Combined together, we propose a complete solution of pre-training a generally reusable model and plugging it as a backbone to accelerate downstream task learning.
[Figure]
Our experiments show that downstream task learning based on a pre-trained embodied set structure can handle unseen changes of task objectives, environmental dynamics and sensor modalities, without assuming any environmental or task distribution priors.
Related work
Pre-training a backbone to accelerate reinforcement learning is not new.
We classify existing works into four categories depending on what is reused in downstream task learning.
(1) Reusable representations : Most existing works fall into this category that aims to learn a representation either to improve sample efficiency or task generalization.
Methods include task-agnostic ones, for example, using VAE (Kingma & Welling, 2013) and a pre-trained ResNet (Shah & Kumar, 2021) for state representations), self-supervised representation learning methods (RAD (Laskin et al., 2020),
CURL (Srinivas et al., 2020) , CPC (Oord et al., 2018), and ATC (Stooke et al., 2021)) for image states, and task-specific ones(Jin et al., 2022), for example, using a behaviour similarity measure (Bi-simulation metrics (Zhang et al., 2020)) to constrain the representation which results in better generalization performance.
(2) Reusable policy/value functions : This approach aims to pre-train a policy or value function that can be further fine-tuned in downstream task learning.
Methods include assuming task distribution priors (meta-RL (Duan et al., 2016)), assuming environmental distribution priors (generalizable RL (Packer et al., 2018) , robust RL (Rajeswaran et al., 2016)), and contextual MDP based ones (Huang et al., 2020).
Recently, representing a policy using generative sequential models has gained increasing popularity (Gato (Reed et al., 2022), Decision Transformers
(Chen et al., 2021)).
However, their capability of downstream task learning is commonly restricted to the training setup, and handling large changes when reused remains challenging (Reed et al., 2022).
(3) Reusable skills and options : This approach assumes tasks can be decomposed into shareable skills (Sutton et al., 2011) that can be learned via time abstraction—learning options as skills.
The difficulty lies in how to discover reusable skills, and it relates to recently proposed works in the unsupervised RL literature (Laskin et al., 2022).
Current methods for skill discovery include curiosity and diversity-based ones (intrinsic controls (Gregor et al., 2016), DIYAN (Eysenbach et al., 2018), behaviour basis-based ones (Proto-value functions (Mahadevan, 2005), and variational intrinsic successor features (Hansen et al., 2019).
Methodology
In the below discussions, we use the Partial Markov Decision Process (POMDP) formation since this setting fits in most real-world reinforcement learning tasks where the full state is not obtainable and the agent is only given observations from its onboard sensors (Igl et al., 2018).
A POMDP is defined as: M =
( X , A , T , r, γ ) , where X denotes the observation space, A is the action space, T is the environment transition dynamics, r is the reward function and γ denotes the discount factor.
3.1 Reusable agent-environment interaction models
(1) Embedded agency: small agent, big world Unlike the Bayesian world model approach (Ha & Schmidhuber, 2018) that aims to construct a big model for the agent as a base knowledge about the world, the Embedded agency (Orseau & Ring, 2012) approach takes the perspective on how the small agent perceives, controls and gradually grows based on its own experience.
Unlike a
Bayesian world model where states and models are objective, an embedded agency approach is subjective(Harutyunyan, 2020) that agent’s past experience will shape its behaviors in new tasks.
This idea of intelligence closely relates to Descarte’s theory of representationalism and mind-body dualism, and to the modern theory of embodied embedded cognition(Haugeland, 1993) in the philosophy of mind research.
The embedded agency approach is appealing in that it promises a scalable learning architecture to build a world model based on the agent’s experience in a self-improvement manner, however, it is also ambitious.
This paper investigates the core idea of embedded agency about how an agent’s past experience can shape its new task learning in a pre-training and reusing paradigm.
(2) On agent-environment boundary Conceptually, from the embedded agency perspective, a generally reusable model in heterogeneous environments and tasks should be the agent itself, since the agent is the only shared component across domains.
Practically, using the agent’s past experience to extract the agent model is difficult since the agent-environment boundary is not the same as we think in a physical system, like a robot.
(Jiang, 2019) provide a detailed analysis of this problem.
(Sutton & Barto, 2018) in chapter 3.1 explains that the agent-environment boundary is often task-specific in different abstraction levels and the boundary could change across tasks and environments, which makes exacting a general agent model from its various task data difficult.
Instead of extracting an agent model, we propose to extract a general agent-environment interaction model that is commonly shared across tasks and environments, and thus can be viewed as an embodiment of the agent.
This proposal is explained below.
(3) Generally reusable agent-environment interaction model In pre-training, suppose the agent’s past experiences D are generated from M different environments and collected using unlimited behaviour policies induced by unlimited reward functions.
This will compose a set of different
POMDPs: {M i } Mi =1 , where M i =
( X i , A , T i , γ ) is a reward-free POMDP that drops the task reward for simplicity.
A without a subscript i means the agent’s action space is unchanged across environments since the embedded agency considers the same agent across environments and tasks.
Let’s assign a one-hot vector encoding Y i as the label for each environment domain M i , then we have a set of domain labels Y
=
{ Y i } Mi =1 , the collected dataset for pre-training can be denoted as
D =
{ D Y i } Mi =1 , where D Y i is the dataset collection in environment domain Y i .
We use successor features (SF)(Barreto et al., 2017) to capture the agent-environment interaction model since SFs summarize the dynamics induced by the environment when following a behaviour policy.
The first step of building a generally reusable agent-environment interaction model is to learn cross-domain transferable successor features in order to extract a general agent-environment interaction model.
For simplicity, we consider a random uniform behaviour policy π 0 for all environment domains.
Recall that, SF is defined as the expected cumulative features of ϕ by following π 0 starting at a specific state.
The SF of ( x i , a ) in environment domain Y i is defined as:
[Equation]
, which satisfies the Bellman Equation as below:
[Equation]
, and it can be learned by minimizing the temporal difference (TD) error
[Equation]
Next, let’s consider a set of domains Y =
{ Y i } Mi =0 .
The goal is to learn a successor feature approximation function f ( . ; θ sf ) :
X i × A → ψ with parameters θ sf that is transferable across all the domains Y .
Inspired by (Feng et al., 2019), we add two constraints to Eq.
in order to make the learned SF cross-domain transferable, using the mutual information definition I ( . ) :
[Equation]
In the first constraint, ψ π 0 is the SF for an arbitrary environment domain, and Y is the corresponding domain label.
It limits the mutual information between an SF and its domain label to a threshold ϵ u in order to make the learned successor feature domain-invariant.
The domain index Y i is dropped since it generally applies to all domains.
The second constraint term maintains the mutual information between an SF and its input domain-specific observations 2 above a threshold ϵ l , in order to prevent the optimization from collapsing to a trivial solution, for example, a random feature space is also domain-invariant but does not include any useful information for task learning.
Technical details about optimizing Eq.
4 are shown in Sec. 3.2.
Note that, learning a cross-domain transferable successor feature representation will not solve the out-of-distribution (O.O.D.) problem when reusing it in downstream tasks with unseen changes, so that directly plugin the pre-trained successor features, as a common approach in ((Barreto et al., 2017, 2020)), will not make our pre-trained model generally reusable .
We propose two techniques to tackle this problem, (1) embodied set construction that discretizes the successor features into prototype sets (Sec. 3.2); (2) feature projection and projected Bellman updates to enable learning stability
-plasticity (Sec. 3.3).
Combined together, we make the pre-trained agent-environment interaction model generally reusable.
3.2 Pre-training: embodied set construction
The pre-training process is proposed as an embodied set construction method that has two steps:
Input: Offline dataset D =
{ D Y i } Mi =1 , ψ π 0 ( x, a ; θ sf ) , embodied set size
N
Result:
Embodied set structure Ω e
Initialize an empty embodied set Ω e = {}
Initialize an empty successor feature vector list L sf = {}
D ← Shuffle ( D ) for each ( x
, a, x ′ ) ∈ D do // Compute cross-domain transferable successor features
L sf ← Append ψ π 0 ( x, a ; θ sf ) end //
Constructing embodied agent state set K-means clustering ( L sf , N ) Ω e = { e i } Ni =1 = cluster-centers as behavior prototypes (1) Learn cross-domain transferable successor features Let a neural network parameterized with θ sf to approximate the SF representation: ψ π 0 ( x, a ; θ sf ) , where ( x, a ) comes from an arbitrary environment domain.
Our aim is to use loss Eq.
4 to find the optimal θ sf .
By adding the Lagrangian multipliers λ u and λ l , the Lagrangian dual of Eq.
4 is:
[Equation]
Directly optimizing the mutual information (MI) terms of Eq.
13 in high-dimensional space is challenging, we provide tractable solutions by approximating the upper bound and lower bound, as detailed in Appendix.
A.1.
Optimizing Eq.
13 will result in transferable ψ π 0 ( x, a ; θ sf ) .
(2) Discretize to behavior prototypes: construct an embodied set structure To facilitate downstream task learning with our proposed embodied feature projection and projected Bellman updates (Sec. 3.3), we discretize the learned cross-domain successor features in three steps: (i) Compute all the successor features using the learned ψ π 0 ( x, a ; θ sf ) from D containing all the environment domain samples.
(ii) Cluster all the successor features using mini-batch K means.
(iii) Collect the center of each cluster to construct a set structure Ω e .
This is summarized in Algorithm 1.
We call Ω e as embodied set since it is the component shared across environments and tasks that can be viewed as the embodiment of the agent.
And each e i ∈
Ω e is called a prototype since it is the center of a cluster that represents a prototype agent-environment interaction behaviour.
3.3 Re-use: a backbone for general downstream task learning
Let a downstream task denoted as M u
= ( X u , A , T u , r u , γ ) , where the observation space X u , dynamics T u , and task objectives r u can be unseen in pre-training.
The action space A remains the same since we assume the same agent learning a new task based on its past experience.
Reusing the pre-trained model under the above heterogeneous settings is difficult due to out-ofdistribution concerns.
Traditional methods will not work.
For example, in methods reusing the pretrained representation (Shah & Kumar, 2021; Kingma & Welling, 2013), changes in the environment will cause non-stationarity in the observation space that directly plugs in the learned representation will make the learning even worse than learning from scratch.
For another example, in methods reusing the pre-trained successor features to fast compute the value functions (Barreto et al., 2020), changes in task objective or the environment will break the linear reward feature assumption.
We propose to avoid directly plugging in the pre-trained model, but to use the pre-trained model— embodied set Ω e as a base structure for projection-based techniques to accelerate the downstream task learning by tackling learning stability and plasticity explicitly.
(1) Stability: Retain previous knowledge by embodied feature projection We define a feature projection operator Π Ω e ( x, a ) , which takes the input of ( x, a ) , localizes its projection on the constructed embodied set Ω e , and returns the localized feature vector e :
[Equation]
find the smallest ξ ( x, a, e )
[Equation]
6) , where ξ ( x, a, e ) =
∥
ψ π 0 ( x, a ; θ sf ) − e ∥ , or it can be any other distance metric functions.
The feature projection operator Π Ω e ( x, a ) retains previous knowledge by always matching the new task experience with the closest prototype e in the embodied set.
Note that Π Ω e ( x, a ) also works for changed sensor modalities.
For example, assuming a task requires the agent to understand textual commands, then adding an extra-textual observation z will augment the observation to [ x, z ] .
The feature projection operator Π Ω e ( x, a ) still applies here by using the unchanged sensory modality part x .
We will show that in Sec. 5.1 (3).
(2) Plasticity: Adapt to changes by projected Bellman Updates We use the below projected
Bellman updates to accelerate learning the unknown downstream task M u =
[Equation]
X u
[Equation]
[Equation]
We maintain two value functions—the task Q π and a projected version V πproj , to support each other’s learning in a bidirectional improvement manner.
The motivation is that the projected function V πproj can learn faster than Q π since it is defined on the pre-trained set Ω e that retains past experience, but is not accurate since Ω e does not adapt to the new task setting.
Meanwhile, the task Q-value function Q π should be more accurate but will take a longer time if learned from scratch.
Learning that alternates Bellman updates Eq.
7 and 8 will play a trade-off between retaining the previous knowledge or adapting to the new tasks.
(3) Reuse example: DQN-embodied While the above two techniques embodied feature projection and projected Bellman updates can generally apply to any RL methods in learning a new task, we use DQN as an example to show how to use them to accelerate downstream task learning.
Assume we use a neural network parameterized with θ u to approximate Q π , and another neural network parameterized with w u to approximate V π proj .
According to Eq. 7 and 8, we compute the target value for Q π and V πproj at training iteration i as:
[Equation]
Then, the learning objectives formulated as an LMSE loss can be written as:
[Equation]
During training, we alternate learning Q π and V πproj .
A full description of the proposed DQNembodied algorithm is summarized in Appendix A.2.
Experimental setup
We aim to test the generality of our method by increasing the difference w.r.t.
the environment, the task objective, and the agent’s sensory modality, between pre-training and reuse.
(1) Choosing the environment
To serve the above evaluation purpose, we require the evaluation environment should be easy to change the task objective (reward), environmental dynamics, and sensor modalities.
We choose MiniGrid (Chevalier-Boisvert et al., 2018) since it is highly re-configurable and easy to reproduce results.
We test tasks including “goal reaching”, “picking up the grey box” and “putting the green key near the yellow ball”.
We use the default partial observability setting that the agent can only observe nearby 7 × 7 cells.
Each cell is described with 3 input values: [OBJECT_IDX, COLOR_IDX, STATE], which gives the object id, cell color and door state (open, closed, locked).
Therefore, the observation x is a 7 × 7 ×dimensional vector.
In changing the sensor modality test, we manually add a textual observation by using the embedding of a text z t to augment x t with [ x t , z t ] .
The agent’s action space remains unchanged during pre-training and downstream task learning, which is a discrete set containing 6 action choices: turn left, turn right, move forward, pick up an object, drop the
object being carried, and toggle (open doors, interact with objects).
(2) Dataset collection
We consider 15 different room layouts environments during pre-training, specifically, 1 layout using 1 room, 4 layouts using 2 rooms and 10 layouts using 4 rooms.
The room layouts are randomly generated.
To fully utilize the agent’s action capability, especially for “picking up an object”, ”dropping an object being carried”, and “toggle”, we randomly place objects (blue keys, blue doors, and different colored balls).
For each room layout, we randomly configure 10 different object placement scenarios.
Therefore, we have 150 environments to generate the agent’s experience
[Figure]
pre-training.
We then let the explore uniformly to navigate the rooms and interact with the objects without any task goals to collect samples.
Therefore, we have 150 task-free POMDPs {M i } 150 i =1 , which result in 150 environment domains { Y i } 150 i =1 , with each Y i is a one-hot encoding vector to denote the domain label.
In summary, during data collection, we created 150 play zones for the agent to explore freely without defining any task goals.
For each domain, we collected 1 M samples which compose the dataset
D =
{D i } 150 i with 150 M samples.
Each samples is a tuple ( x, a, x ′ ) .
[Figure]
Following (Hoang et al., 2021), we pre-train a variational auto-encoder (VAE (Kingma & Welling, 2013)) that maps observation x t to a feature vector ϕ t , which is then used to learn the cross-domain transferable successor features by learning the function ψ π 0 ( x, a ; θ sf ) using methods proposed in Sec. 3.2.
During embodied set construction, we set N = 10 K to discretizes all SFs into 10K clusters and construct the embodied set Ω e representing 10K prototypes behaviors.
(4)Re-use: a backbone for downstream task learning
Given the pre-trained embodied set structure Ω e , we test how it can be reused to accelerate downstream task learning using our proposed
DQN-embodied method (see Algorithm 2), as described below.
Preliminary results
We conduct the generality test by adding differences between pre-training and reuse.
Fig.
3 gives an overview of task setup in the below three tests.
Unseen task objective + unseen environment + unseen sensor modalities
We test the generality of our proposed method by adding an unseen task objective, environment settings, and sensor modality in downstream task learning.
The task is “putting the green key near the yellow ball” wherein the agent needs to understand “green key” (unseen in pre-training), “yellow ball” and “near”.
The task requires the agent to find the green key, carry it and drop the key in a cell nearby the yellow ball.
The reward function is defined as r ( x, a, x ′ )
= 0 .
1 (if pick up the green key), 10.0 (if succeed).
Adding textual sensor modality: Note that this task requires text prompts otherwise is difficult to solve.
We manually design three types of textual observations for the agent— The mission statement “Put the green key near the yellow ball” + different textual statuses: {“and I have nothing”, “and I have the [color] [object]”} depends on what the agent is carrying.
The textual observation it then converted to an embedding vector o using a pre-trained sentence transformer MiniLM (384 dimensional output, 80 MB model size) (Wang et al., 2020) using a libary provided by (Reimers & Gurevych, 2019).
Therefore, the agent’s observation at time t is augmented as : [ x t , o t ] , where x t ∈ R 7 × 7 × 3 is the same observation as in pre-training, and o t ∈ R 384 is the extra sensor modality added to the agent to provide textual observations.
Baselines: We compare our method ( DQN-embodied ) to the following baselines.
(1) Without pretraining: vanilla DQN (Mnih et al., downstream task learning; (4) Without textual observations (DQN-no-text) using a vanilla DQN (Mnih et al., 2013).
Note that baselines DQN-VAE can not directly handle an extra sensor modality.
When taking the o t as inputs to the above baselines, we concatenate the VAE encoded observation with o t .
We test the downstream task learning performance of all methods by evaluating the success rate every 10K training steps using 100 independent runs with different random seeds.
Results (Fig. 4) show that (1) Without textual observations, a vanilla DQN agent can not solve the task; (2) Since the new task setting is heterogeneous compared to tasks in the training pools, directly reusing previous experience (the ER baseline) will deteriorate the downstream task learning performance (vanilla DQN is better than ER); (3) Also, since the new task is never seen during pre-training, DQN based on pre-trained VAE performs the worst.
(4) For the test task that has unseen task objectives (put a color object near another color object), unseen environmental settings (box object never seen in pre-training), and unseen sensor modalities (adding textual observations), projecting the downstream task samples on the pre-trained embodied set will result in better sample efficiency than all baselines.
Conclusion
This paper tackles the problem of pre-training a generally reusable model for downstream tasks with unseen task objectives, environments, and even with unseen sensor modalities.
Our approach is inspired by the embedded agency perspective that the agent’s past experience shapes its future decision-making behaviors.
We propose to learn an agent-environment interaction model by extracting cross-domain transferable successor features using the agent’s various task experiences.
To make the model generally reusable, we propose to discrete it into behavior prototypes, resulting in an embodied set structure that further supports embodied feature projection and projected Bellman updates to enable downstream task learning with heterogeneous settings.
Preliminary results show that downstream task learning by projecting on our proposed embodied set structure can handle unseen changes in task objectives, environmental dynamics and sensor modalities.
For future works, more comprehensive experiments are needed to understand why and how the pretrained embodied set structure can be generally reused in downstream task learning.
It’s also worth working on methods that can continually update the constructed embodied set structure given that more task experience is gained, thus forming an embedded agency version of continual reinforcement learning (Khetarpal et al., 2020) solution.
A Appendix A.1 Learn cross-domain transferable successor features
Assume a neural network with parameters θ sf is used to approximate the SF representation denoted as ψ π 0 ( x, a ; θ sf ) , where ( x, a ) comes from an arbitrary environment domain.
Our aim is to use loss
Eq.
4 to find the optimal θ sf .
By adding the Lagrangian multipliers λ u and λ l , the Lagrangian dual of Eq.
4 is:
[Equation]
Directly optimizing the mutual information (MI) terms of Eq.
13 in high-dimensional space is challenging, we provide tractable solutions by approximating the upper bound and lower bound.
Upper bound of I ( ψ π 0 , Y )
Following (Feng et al., 2019), the I ( ψ π 0 , Y ) is upper bounded by replacing one of the marginal distributions with a variational posterio distribution.
For any distribution q ( Y ) , we can have an upper bound of I ( ψ π 0 , Y ) :
[Equation]
, where q ( Y ) can be empirically estimated in a non-parametric way, like using kernel density estimation from the dataset.
Therefore, minimization of the KL-divergence term will result in minimizing the mutual information term, which can be formulated as an adversarial training objective, as derived in (Feng et al., 2019):
[Equation]
(15) , where q θ u
( Y | ψ π 0 ) is a classifier that predicts the probability of the sample belonging to label Y , and q ( Y ) is a constant that can be dropped during optimization.
Lower bound of I ( ψ π 0 ,i , x i )
There are multiple ways ( (Oord et al., 2018; Nowozin et al., 2016)) to approximate a lower bound of the mutual information term.
We follow (Feng et al., 2019) that maximizes the MI term by maximizing the Jensen-Shannon divergence form of MI, ˆ I JSD ( ψ π 0 ,i , x i ; θ l ) , where θ l parameterizes the JSD.
Readers can refer to (Nowozin et al., 2016;
Feng et al., 2019) for more details.
A complete loss term of Eq.
13 can be rewritten as a minimax optimization objective:
[Equation]
, which is now tractable.
A.2 DQN-embodied
By using the projected Bellman updates (eq. 7 and eq.
8) proposed here, we compute the target values as described in eq. 9 and eq. Then, the full DQN-embodied algorithm alternates optimizing Q π and V πproj is as below:
Given pre trained embodied set Ω , feature projection operator Π Ω ( x, a )
Initialize Q u
V
; w u ) and replay buffer D for i
N do // do Q
Set y
i x a x ′ )
∼T u [ r u + γV
Π Ω e x ′ max a ′ x ′ a ′
θ u i − w u
−
Perform
Set
Perform end end