Leveraging Mutual Information for Asymmetric Learning under Partial Observability
Authors: Hai Huu Nguyen
Long Dinh Van The
Christopher Amato
Robert Platt
Venue: CoRL 2024 (2024)
Keywords: Partial Observability
Mutual Information
Reinforcement Learning
TL;DR: This paper proposes to use the mutual information between the history and the state to improve learning under partial observability, assuming state availability during training.

Abstract: Even though partial observability is prevalent in robotics, most reinforcement learning studies avoid it due to the difficulty of learning a policy that can efficiently memorize past events and seek information. Fortunately, in many cases, learning can be done in an asymmetric setting where states are available during training but not during execution. Prior studies often leverage the state to indirectly influence the training of a history-based actor (actor-critic methods) or a history-based critic (value-based methods). Instead, we propose using state-observation and state-history mutual information to improve the agent's architecture and ability to seek information and memorize efficiently through intrinsic rewards and an auxiliary task. Our method outperforms strong baselines through extensive experiments and achieves successful sim-to-real transfers to a real robot.

Leveraging Mutual Information for Asymmetric Learning under Partial Observability
Hai Nguyen', Long Dinh², Robert Platt', Christopher Amato!
Khoury College of Computer Sciences, Northeastern University, Boston, MA, United States
Horus AI, Hanoi, Vietnam
https://sites. wwwp:wiweooooiaaaaeeeeeeeeee pomp-
Even though partial observability is prevalent in roboties, most rein- forcement learning studies avoid due to the difficulty of learning policy that can efficiently memorize past events and seek information. Fortunately, in many cases, learning can be done in an asymmetric setting where states are available during training but not during execution. Prior studies often leverage the state to indirectly influence the training of a history-based actor (actor-critic methods) or history-based critic walue-based methods). Instead, we propose using state- observation and state-history mutual information improve the agent's architec- ture and ability seek information and memorize efficiently through intrinsic re- wards and an auxiliary task. Our method outperforms strong baselines through ex- tensive experiments and achieves successful sim-to-real transfers to real robot.
Abstract:
Keywords: Partial Ohservability, Mutual Information, Asymmetric Learning
Introduction
While partial observability is ubiquitous in robotics [1, 2], most reinforcement learning (RL) robotics research avoids it due the difficulty of training policy that can effectively memorize past infor- mation and perform information toformationggathrinng actions. Fortunately, full near-full observability dur- ing training is possible in many real-world domains using simulator model. In the literature, this is the asymmetric learning setting [3, 4, 5, 6, 7, 8], in which privileged information such as states can be available during training. Previous work in this setting often learns a state- based critic to indirectly influence the learning process of a) a history- -based actor [3, 9] using actor-critic RL to output actions directly, or b) a history- critic [6] value-based RL to extract actions from. In contrast, using an information-theoe intrinsic rewards and optimize an auxiliary task, shaping the agent's behavior directly.
generative environment
approach, we instead use the state generate
Specifically, this study proposes to use the mutual information between the state and the history during training to motivate active information gathering and to improve task-relevant memorization. Our intuition is that active information gathering mainly involves discovering actions that help the agent reveal more about the environment's state, gradually moving from partial to full observability. Therefore, the amount of mutual information gained per step can be useful intrinsic reward to encourage informative actions. Moreover, we train the agent predict the state features from past histories as auxiliary task to help the agent focus on memorizing enlightening past events, i.e., previous actions or observations that have gained more information about the environment state.
We propose a two-stage approach to turn intuition into a practical and efficient method. The first stage extracts task-relevant but non-overlapping state and observation features by performing super- vised learning using a state-contained dataset. Specifically, we predict the reward and the dynamics in the feature space while minimizing the mutual information between state and observation features. We perform task learning in the second stage, which involves training a transformer-basee sequence model to output a history summary with maximal mutual information with the learned state features
8th Conference Robot Learning (CoRL. 2024), Munich, Germany.

---

as possible. The goal is to reduce the memory burden of the sequence model by encouraging it to avoid memorizing readily available observation features. With a purposefully designed structure, intrinsic rewards, and the auxiliary task, our approach outperforms other strong baselines in exper- iments on six domains. Moreover, we carry out successful sim-to-real transfers to verify that the policies learned in simulation can perform well in three robot manipulation tasks.
2
Preliminaries
We first discuss the partially observable Markov decision process [10, 11] (POMDP) framework, methods that minimize/maximize the mutual information, and describe our problem statement.
2.1 Partially Observable Markov Decision Process
POMDP defined by a tuple (S,A,9, R. O). where S. and Q the state space, the action space, and the observation space, respectively. In POMDP, the transition function ['ss, a, s') gov- ems state changes, and the observation function Daa,'', controls how observations are emitted. As the current observation only partly reflects the state, acting optimally in a POMDP with a high level of partial observability often requires information gathering and a memory-based policy that selects actions based on the entire cesionnobservation history h, (oo, .p...., -, 0,) [12].
2.2 Minimizing and Maximizing Mutual Information
Mutual Information (MI) measures the dependence between two random variables and y:
p(x,
(x:; v) Ep(x.v) log p()))(y)
where p(x,y) y) is the joint distribution and p(x)), p(y) are the marginals.
Minimizing MI:
Cheng et al. [13] proposed to minimize [xx; v) by minimizing upper bound named Contrastive Log-ratio Upper Bound (CLUB), which is formulated as:
ICLUB == Bpt [log p(y] [x]] Ep En(o) [lo [[][=)]] 2 I(x; v)
(2)
In practice, p(y(z) is often unknown and approximated by learnable variational distribution g(y) (z). Maximizing MI: Instead of maximizing 7 (<r, v), Deep InfoMax (DIM) [14] is a method that max- imizes the MI between x and a high-level representation 2² of x. The process involves an encoder E: X' + n" and a classifier : N' x 2 + IR to discriminate between samples from p(x,=*) and the product of marginals p==)p(=*). DIM uses the Jensen-Shannon MI estimator [15]:
IIM == E((), (= (xx [(x))) Eegee,ee E(x))) E(x)) < B(x)
(3)
where sp(z) = log(1 + e²) is the softplus function. Maximizing [(z, 2²) involves continuously
updating E and by maximizing the estimator JDIM through gradient ascent.
In practice we optimize Eq. (2) and Eq. (3) through samples; see the Appendix for more details.
2.3 Problem Statement
We focus on the asymmetric RL setting [3, 9, 6, 8] in which, during training, both the state and the observation o, are available while only is accessible during execution. Unlike an MDP, the goal here is to find history-based (not reactive) policy <aa (1)) which maximizes the expected discounted return J = E Eie x² R(st, at)]. where E [o. 1) is discounting factor. We use an information-base based approach to leverage states during training for efficiently learning π (h).
Method
Our method has two stages. First, we learn non-overlapping state and observation features from fixed dataset containing privileged state information. The dataset is generated using random agent

---

or a simple hard-coded planner. Next, we use the learned features to train memory based agent with a two-branch architecture using information-based intrinsic rewards and an auxiliary task.
3.1 Learning Task-Relevant and Non-Overlapping State and Observation Features
This stage aims to learn state features =* (encoded by a state encoder Φ) and observation features z" (encoded by an observation encoder v)) with two desired properties. (P1): Firstly, concatenated feature 2d z" to encode task-relevant state features. In other words, 2" + 2" is compact "state", which is useful when the raw state is high-dimensional or contains task-irrelevant features. (P2): Secondly, we want =" to contain unobservable task-relevant features, while x" con- tains observable tusk-relevant features. With Pl and P2 satisfied, 2 will encode only task-relevant but unobservable features inferred from the history as =" already contains features extracted from the observation. We can also view 2 as the outcome of an ideal information gathering process, which results in task-relevant hidden information not captured by the current observation (i.e., =9).
the
We satisfy (PI) by training z" + =" given an action to be predictive of 2"", zd and the reward similar to bi-simulation studies for MDPs [16, 17, 18J. For (P2), Wang et al. [8] minimized the KL divergence between d(s) and a fixed standard normal distribution by giving a penalty whenever features are derived from the 8 instead of o. However, we will empirically show (in Section 5.2.1) that this approach does not satisfy (P2). Instead, inspired by [19], we directly minimize the mutual information [ 2²) using Eq. (2). However, minimizing only [[2"; (() may result in undesired scenario, where [2";22" 2") small but only 6(s) or (0()) captures all good features, and the remaining contains no useful information. Therefore, we further regularize learning with maximizing I(o; 2²) and (8) using Eq. (3). Algorithm summarizes this stage (more details in the Appendix C).
Algorithm Learning 2 and 2" (more details in Appendix C)
Require: Dataset D {(s, o,, a, Ti, o
1: repeat
2:
Sample a batch of data {(5 0z a, ri, s ) ~
3: Compute current features 2f o(si), 2 (0:)
Compute reward and next features using a dynamics model g: fr, 28, 2. = g(25 2,,
4:
5:
Compute target features z²¹ o(5)), 2 = ((f))
6:
Compute L, = (r+ - )² L. (stop-grad(2;) Compute LCLUB using Eq. (2) to minimize 2²²";²")
2, 113, C. ||stop- 2
7:
8:
Compute LDIM using Eq. (3) maximize / (0; 2") and =s);2")
9:
Update o, vi, g to minimize X, + X,C, + AoCo - ACLUB CCLUB + ADIM until convergence
3.2 Learning Partially Observable Tasks
After learning 2" and 2", we now introduce memory-based agent with two-branch architecture (see Fig. 1) for task learning using RL. The memory-based top branch is responsible for encoding the history using a GPT-v2-based sequence model [20, 21, 22). In contrast, the memoryless bottom branch encodes the current observation. When the agent finishes gathering information, the top branch will ideally output features in 2", which are not captured in 2a. For the bottom branch, we directly (o) initialize it. A two-branch architecture is new by itself (see [23, 24, 25, 26]) and has been found to offer no performance gain over the single-branch (no bottom branch) architecture [23]. Unlike the previous work that used this architecture, by having non-overlapping =" and N and using z" to bias the top branch's 's learning, we alleviate the top branch's 's memory burden by training it to focus only on features not readily observed.
Because z" is the ideal outcome of the information-gatherin process, we encourage information- gathering by rewarding actions that produce h2 that maximize I[h; 2f). Moreover, we use maximizing [(h; 2f) an auxiliary task to help memorize important information from history.
3

---

Concet
Concet
Intrinssic Reward for Information Serking
Top Branch Predict o' from
O2
Encoder
+
Sequence Model
+
Vhha)
total pp + or intrinst:
Encoder
pralciiaic xr I(hy; 1(
Bottom Branch Init by z"
vi(o)
Auxilary Loss for Memorizing Useful Information
er CRu + B]]= k(h))l
Figure 1: Our proposed agent structure with two branches ( denotes vector concatenation). The observation encoder (b(o) is directly used for the bottom branch. An intrinsic reward calculated based (h, =*) encourages information seeking. Moreover, to better memorize useful informa- tion, besides being trained to minimize the RL losses, the top branch (k) is trained to predict the state feature 2² from the history he.
Encouraging Information Seeking. By maximizing I [hh;;22) the agent is encouraged to gather information by acting to get he that encodes maximal information with 2f- We closely follow [27] and maximize the mutual information by maximizing the following lower bound [28]:
T(ba; =) = (= - H (2 2 (=() + Eny Dlog fo(=ā)
t-1
log fo(z;
= (2 + log fa (ho + EN
(4)
where [o(=*[h) serves as an approximation to the true conditional distribution p(=*(h). Similar to [27], we assume Ih(=*[h) follows a Gaussian distribution centered around deterministic history
embedding so(h) with variance I²I. We parameterize fo sh) using the GPT-based sequence model and train fo (h) to predict 2 across batches of (h, 2*) samples by minimizing:
Dj = (h,
[||=* fo(h)
(5)
Since only the third term of Eq. (4) depends on π. we interpret log lole To(s he-1] as an intrinsic reward term of pintrinsic This intrinsic reward captures per-step incremental information gain that the current history h4 reveals about z compared to the previous h4-1- With a weighting factor our final reward therefore becomes r == rta + cor-ittrinsic, In practice, we calculate the intrinsic rewards as deusine (2f, fo deasine (2i fo (h), with deusine being the cosine similarity distance.
rt piitrinsicc
Encouraging Information Retention. While the intrinsic reward encourages information gather- ing, we also want the agent to memorize important observations (or actions). For this purpose, train the agent's top branch (k in Fig. 1) to predict =² over a batch of episodes as auxiliary task using the same Eq. (5) but replacing f with k. While f and k seem similar, the key difference is that k is part of the agent (it is the top branch) while f is not (f is a network that approximates p(=*lh)). Training the agent this way can be considered an indirect and "online" way to maximize 2(h;2") that the agent can mentally memorize important events within the history that increased [(h; =*).
4 Related Work
Leveraging privileged information during policy learning has been successful in various settings, most notably in the multi-agent setting. Usually, joint quantities (e.g., observations, histories [29], or actions) from all agents, the state [30], or their combinations [31, 33, 34], are used to train a centralized critic to aid training decentralized actors. In the single-agent setting, Pinto et al. [3] leveraged the state availability in simulation to improve an image-based policy. Under partial ob- servability. Belief-grounder Networks [5] use an auxiliary loss to reconstruct the privileged state distribution information from history representations. Using a value-based method, Baisero et al. [6] uses the state access to learn history-based critic. Using the actor-critic framework, Baisero and Amato [9] learn a history based critic to help learn history-based actor. Concurrent with our work is [35], which also proposes informational rewards to encourage information gathering.

---

Bool
Guel
Goad
Goal Padi
Unlocked
Immowable
Immovable
Movabie
Movable
(D-P) Sptinx
(0-F) CarFlag-20
(C.P) Pullin
Block-Pusting
(CP) Lrawer-Opening
Figure Domains. D/C: discrete/continuous actions. V/P: cestor-basedppixllll based observations.
While these methods use the privileged information additional input (e.g., .. train critic), first find compact and task-relevant representation of the privileged information (similar [8]). However, diverge from [8] by directly using privileged information for encouraging informa- tion gathering. Our information- based approach shares similarities with meta- meta-learning work [27], which uses the mutual information between the history and the training task ID to compute intrinsic rewards to encourage information gathering. While [27] learns two policies for exploration (with the intrinsic rewards) and exploitation (with the environment rewards), we learn single policy using a combined reward and add an auxiliary loss to improve memorization. Our idea of minimizing the overlapping between = and =" has connections with the mixed-observability framework [36], which removed the overlapping between state and observation to reduce the computation for state distri- bution tracking. However, the proposed method [36] requires knowing the complete dynamics (i.e., planning- based) and the state and observation structure, while ours assumes no such knowledge.
Simulation Experiments
We perform experiments on three grid-world domains with discrete action spaces and three robot manipulation domains with continuous action spaces (see more details in Appendix B).
5.0.1 Discrete Action Space Domains
In Sphinx [8], CarFlag-2D [37], and Heaven-Hell [6] (see Fig. 2), agent must reach randomly generated goal position (m). Specifically, the goal can be one of the green corners in Sphinx, left/rihht ends in Heaven-Hell, and any cell in CarFlag-22 2D. Normally, the observation only contains the agent's location. Only when entering the information region( the agent can additionally observe the goal. For instance, when entering blue cell in Heaven-Hell, the side of the goal (left/right) is observed. Similarly, in the Sphinx and CarFlag-22 2D, the observation will encode the goal cell if entering the blue region. Rewards are given sparsely: +1 when the agent reaches the goal, -1 when it reaches the trap locations (m) if any, and otherwise. In Sphinx, an additional penalty of -0.2 is given when entering the blue cell, making the iiformatioongatheriii process costly. In these domains, the state encodes both the current observation and the goal cell.
5.0.2 Continuous Action Space Domains
In Block-Pulling, Block-Puste Puscing and Draser-Opening (modified from domains in [38]). robot must pull, push, or open the only movable object (a block or a drawer) among two, which appear the same under top-down depth camera. The challenge is that the mobility of the objects is hidden if only relying on the current depth image (the observation). An optimal agent must perform information information-gathering gathering actions to determine the objects' mobility and then memorize the information while manipulating the movable object to achieve the tasks. Here, the state is also pixel-based, created by concatenating the observation and a 1-channel image that masks everything except the movable object. The reward is +1 only when finishing the tasks. To overcome the reward sparsity, we give 80 episodes of demonstrations pre-populate the replay buffers of all agents. Moreover, because these tasks exhibit spatial symmetry [39], applied consistently for all observations and actions within episode, for data augmentation.
random SO(2) rotations (see Appendix F),

---

5.1 Baselines
We use DDQN [40] (discrete action spaces) and SAC [41] (continuous action spaces) as the base RI. algorithms for our agent and all baselines. We make them memory-based using a simplified version of GPT-v2 [20, 22] as the sequence model. All following baselines use the two-branch architecture like ours (see Appendix A).
DTQN, TSAC: These variants [21, 22] are based on DQN and SAC and use transformer (a simpli- fied version of GPT-v2 (20]) as the sequence model. ZP-DRQN, ZP-RSAC: Ni et al. [42] proposed the next hidden state prediction as an unsupervised auxiliary task to improve learning under partial observability. We applied their method to DRQN and RSAC and used a GRU [43] instead of trans- former here to satisfy the recurrent encoder property in [42] (i.e., a transformer is not a recurrent model). BA-DTQN, BA-TSAC: The biased asymmetric version of DTQN and TSAC. For BA- DTQN, (s, and (h,, a) are learned during training, and argmax, Q(h, a) is the selected action during deployment. For BA-TSAC, only the state-based critic Q(s, a) is learned to help learn history-based actor. These asymmetric agents can sometimes perform well [3] but were found [6, 9] to introduce learning bias, i.e., Q(8,a) a) is a biased estimate of Q(h, a). UA-DTQN, UA-TSAC: We implement unbiased asymmetric variants of DTQN and TSAC using the framework proposed in [6, 9] (originally applied to DRQN [44] and A2C [45]). To fix the bias issue, during training, these methods learn Q(h, 8, (instead of (s, a)) and Q(h,a) and use Q(h, a) to pick actions dur- ing execution. B-DQN, B-SAC: Believer [8] originally trains GRU-based history encoder, similar to our first stage. For fair comparison, we use the same GPT replace the GRU and the same data set for our agent and these agents. Because the encoders already encode the history, their outputs are directly used for memoryless RI. algorithms, i.e., DQN or SAC in our case. Furthermore, we fine-tune the history encoders to strengthen these baselines, following the public code.
3.0
DTON
DION BIOO DOO (WNON
DTON
1.0
OTTN
0.6
0.8
0.75
0.DCK 24-ORDN
DRON
Rutumn
Qura
0.6
Chr's
Our
0.2
0.6
0.25
0.2
-022
-0.25
-0.4
0.0
-0.50
100k
2002
B00k
400%
200k
800K
1.0M
100k
2002
300k
400A
500K
Emviromment Step
Environment Step
Er Environment Step
(a) Sphinx
(b) CarFlag- 2D
Heaven -Hell
1.0
BA-TAAC
BA-TSAC
BA-TAC UATTAC
0.8
UA-TSAC TSAC
0.8
LATSAC
0.81
Batum
Retum
0.6
0.6
0.2
0.7
0.0
0.0
BCR
1.00
Envionment
Environmenc Stop
Environmment
(d) Block-Pulling
(c) Block--uushing
(f) Draver-Opening
Figure 3: Learning curves averaged over five random seeds with shaded mean and one std. Dashed vertical lines denote when the task learning stage of our method and B-DON/B-SAC starts.
5.2 Experimental Results
As shown in Fig. 3, our method is generally superior to all baselines in all domains. Among baselines with state access like ours, only unbiased asymmetric methods (UA-DTQN/-TSAC) can compete with ours in Heaven-Hell, based baselines (B-DQN/-SAC) do not perform well in all domains despite our significant effort to strengthen them (replacing GRU with GPT, hyper-paramen tuning, and online fine-tuning). When information seeking is hard (c.g.. in Sphinx with clarifying penalty, in Heaven-Hell with
and Drawer-Opening, Most surprisingly, Believer-

---

1.0
0.8
0.0
Returs
(Dun))
0.6
onty (Dam)
o(s)
Concal
Trisk Learning
0.4
0.0
0.0
*((0)
-02
-0.2
100%
200V
300k
400K
500k
100k
200k
300k
400K
500K
Environment Stop
Envfronment Stop
(a) Test structure
(b) Ours
(c) Believer [8]
Figure 4: a) Overlapping test: using 2 or concatenating 2" and 2" for task learning: b-c) Learning curves in Heaven-Hell when using our method (Section 3.1) or Believer [8] to learn 2 and 29.
goal-avoided path for goal information), only our agent can perform well by using the intrinsic rewards. When the information seeking is easier, our method still outperforms (e.g., in CarFlag-22 Block-Pulling, and Drawer-Opening) thanks to pre-trained compact features.
5.2.1 Learned High Quality and Non-overlapping Representations
To show that our method can achieve both (P1) and (P2) in Section 3.1, we perform task learning in Heaven-Hee using a) 2 2", b) only 2², and c) raw state s. We compare two ways to learn the representation:
method and Believer [8]. From Fig. 4, x 22 yields better performance than the raw state. This indicates that our method can learn high-quality state features. As expected, using 2" alone fails to solve the task because 2" will only contain the features of the goal's position. The task can only be solved with the aid of 2", which contains the agent's position. In contrast, when using Believer [8] to learn 2 and 2", using either 2 @ 2 or 2 makes no difference: both can solve the task, indicating that 2 also contains the agent's position, which is already contained in 2
5.2.2 Ablation Studies
In this section, perform ablation studies to analyze key design choices: using GPT instead of recurrent neural network (e.g., GRU) as the sequence model and b) using both the auxiliary losses and the intrinsic rewards. We only report representative results in this section; see the Ap- pendix D for more details.
1.0
1.0
0.8
0.8
0.8
0.6
0.6
0.4
CRU
0.4
0.7
0.0
0.0
0.0
200k
6O0k
1.0M
Environmont Stop
Enviromment Stop
Environment Stop
(a) CarFlag-2D
(b) Block-Pulling
(c) Blloc-Puushing
Figure 5: Comparing when using GPT or GRU as the sequence model in our approach.
Using GPT instead recurrent neural network. While a GRU [43] might have trouble memo- rizing important information long sequences (i.e., saw important observations but forgot), GPT attend and memorize any past information. This property is crucial for our agent maximize I [hh; ((s)) effectively. From the performance gap in Fig. 5, we can see the important role of GPT.
Using both auxiliary losses and intrinsic rewards. While the auxiliary losses help information re- tention, intrinsic rewards are needed to encourage information gathering directly. As shown in Fig. 6, in domains when information gathering is not challenging, such as in CarFlag-2D (the information region quite large) and the robot domains (agents are provided with demonstration episodes), ing auxiliary loss is sufficient for good performance. In Sphinx, when information is harder to seek because of the involved cost, the intrinsic rewards are essential to improve the performance.

---

1.0
0.75
0.8
Ausil boalaayyttttnet
0.6
Inttralc Only
0.6
0.00
Bry
0.4
-025
0.2
0.2
-0.75
Acallar-Or5 Onl
0.0
0.0
100k
2001
300k
403k
500V
200R
40.x
600R
800k
1.0M
200
40k
60k
800
Enstronment Step
Environment Step
Enstronment Step
Sphinx
(b) CarFlag-2D
(c) BlockockPPP Bllck-Puushing
Figure 6: Comparing using intrinsic rewards or auxiliary losses versus using both.
6 Sim-To-Res Transfers
Cancra
Domain
Ours
UA-TSAC
Block-Pulling 45/50
34/50
Block-Pushing
44/50
12/50
Drawer-Opening 41/50
32/50
Table I:
Average success rates of sim-to-real transfers between our method and UA-TSAC (best baseline).
Figure 7: Real robot setup.
To verify the performance of the learned policy in simulation, we perform sim-to-real transfers using the best policies in simulation UR5 robot (see Fig. 7). To best obtain non-occluded top- down depth images, we combine the point clouds from two side-view cameras and one top-down camera and project at the gripper's position (see Appendix E). Furthermore, during training, we add Perlin noise [46] to the observation to reduce the gap between the simulated and the real-world depth images. We perform 50 evaluation episodes, divided equally into two cases when the policies first manipulate the immovable or movable objects. Results from Table I show that the learned policies can be zero-shot transferred in the real world (see our supplementary video for more details), and method performs better than the best baseline (UA-TSAC). During sim-to-real transfers, the main failure mode is when the robot goes down to perform actions; collisions with objects (e.g., blocks, drawers) might happen. This can trigger the robot's protective stop and terminate the episode. Additionally, in Draver-Opening, the transferred policies sometimes clumsily move one drawer far away from the other, creating a novel scene never seen in simulation.
Conclusion and Limitations
This work proposed an information based method that directly uses the mutual infor- mation between state and history to encourage cearmaationgatherriie under partial observability. We also tackle memorization by incorporating an auxiliary task to predict the state feature from the history alleviate the agent's memory burden by only requiring its memory component to memo- rize unobservable task-relevant features. Experiments show that our method is effective and scalable to domains with high-dimensional states and observations with deployable working policies.
Conclusion.
Limitations. Our method's s biggest limitation the challenge learning good and non-overlapping representations at the first stage. practice, we need to select the right dataset to train with. We random interactions the grid-world domains, but expert demonstrations training in the robot domains that the representations leamedd representations can satisfy both (P1) and (P2), need to iterate between hyper-paramete tuning and performing the overlapping test in Section 5.2.1. We also want to note that while the second stage works optimally with high-quality representations, our method remains effective with less ideal features, as the agent can still refine these through reinforcement learning.
needed for useful. In practice, guarantee the

---

Acknowledgments
This research is supported by grants from the Army Research Office (W911NF20-1-0265)) NASA (8445 and the National Science Foundation (NSF 2024790, 1750649, 2107256, 2314182, 2409351).
References
[1] H. Kurniawati. Partially observable markov decision processes and roboties. Annual Review of Control, Robotics, and Autonomous Systems, 5:253-277, 2022.
[2] M. Lauri, D. Hsu, and J. Pajarinen. Partially observable markov decision processes in robotics: survey. IEEE Transactions on Robotics, 39(1):21-40, 2022.
[3] L. Pinto, M. Andrychowicz, P. Welinder, W. Zaremba, and P. Abbeel. Asymmetrie actor critic for maage-aased robot learning. RSS, 2018.
[4] D. Chen, B. Zhou, Koltun, and Krähenbühh. Learning by cheating. In Conference on Robot Learning, pages 66-75. PMLR, 2020.
[5] H. Nguyen, B. Daley, X. Song, C. Amato, and R. Platt. Belief-grounde networks for accel- erated robot learning under partial observability. In Proceedings of the 2020 Conference on Robot Learning, volume 155, pages 1640-1653. PMLR, 2021.
[6] A. Baisero, Daley, and C. Amato. Asymmetric dqn for partially observable reinforcement learning. In The 38th Conference Uncertainty in Artificial Intelligence, 2022.
[7] H. Nguyen, A. Baisero, D. Wang, C. Amato, and R. Platt. Leveraging fully observable policies for learning under partial observability. In 6th Annual Conference on Robot Learning, 2022.
[8] A. Wang, A. C. Li, T. Q Klassen, R. T. Icarte and S. A. Mcllraith. Learning belief representa- tions for partially observable deep rl. In International Conference on Machine Learning, pages 35970-35988. PMLR, 2023.
[9] A. Baisero and C. Amato. Unbiased asymmetric reinforcement learning under partial ob- servability. In Proceedings of the 21st International Conference Autonomous Agents and Multiagent Systems, pages 44-52, 2022.
[10] K. J. Aström. Optimal control markov processes with incomplete state information Jour- nal of mathematical analysis and applications, 10:174-205, 1965.
[11] L. P. Kaelbling, M. L. Littman, and A. R. Cassandra. Planning and acting in partially observ- able stochastic domains. Artificial intelligence, 101(1-2):99-34, 1998.
[12] P. Singh, Jaakkola, and Jordan. Learning without state-estimation partially observ- able markovian decision processes. In Machine Learning Proceedings 1994, pages 284-292. Elsevier, 1994.
[13] P. Cheng, W. Hao, S. Dai, J. Liu, Z. Gan, and L. Carin. Club: A contrastive log-ratio upper bound of mutual information. International conference on machine learning, pages 1779- 1788. PMLR, 2020.
[14] R. D. Hjelm, A. Fedorov, S. Lavoie-Marchildon, K. Grewal, P. Bachman, A. Trischler, and Y. Bengio. Learning deep representations by mutual information estimation and maximization. arXiv preprint 1808, 2018.
[15] S. Nowozin, Cseke, and R. Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. Advances in neural information processing systems, 29, 2016.

---

[16] A. Zhang, R. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006. 10742, 2020.
[17] M. Kemertas and T. Aumentado-Armstrong Towards robust bisimulation metric learning. Advances in Neural Information Processing Systems, 34:4764 2021.
[18] C. Gelada, S. Kumar, J. Buckman, O. Nachum, and M. G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International conference on machine learning, pages 2170-2179, PMLR, 2019.
[19] L. Zhao, Y. Wang, J. Zhao, L. Yuan, J. J. Sun, F. Schroff, H. Adam, X. Peng. D. Metaxas, and T. Liu. Learning view-disentangled human pose representation by contrastive cross-view mutual information maximization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 77793-12802. 2021.
[20] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[21] K. Esslinger, R. Platt, and C. Amato. Deep transformer q-networks for partially observable reinforcement learning. arXiv preprint arXiv:2206601078, 2022.
[22] T. Ni, M. Ma, B. Eysenbach, and P.L. Bacon. When do transformers shine in rI? decoupling memory from credit assignment. Advances in Neural Information Processing Systems, 36, 2024.
[23] Ni, B. Eysenbuch, and R. Salakhutdinov. Recurrent model-free RL can strong baseline for many POMDPs. In International Conference on Machine Learning, pages 16691-16723. PMLR, 2022.
[24] D. Ha and J. Schmidhuber. World models. arXiv preprint arXiv:
2018.
[25] L. Zintgraf, K. Shiarlis, M. Igl, S. Schulze, Y. Gal, K. Hofmann, and S. Whiteson. Varibad: a very good method for bayes-adaptive deep rl via meta-learning In International Conference on Learning Representations (ICLR), 2020.
[26] C.-C. Hung, T. Lillicrap. J. Abramson. Y. Wu, M. Mirza, F. Carnevale, A. Ahuja, and G. Wayne. Optimizing agent behavior over long time scales by transporting value. Nature communications, 10(1):1-12, 2019.
[27] E. Liu, A. Raghunathan, Liang, and C. Finn. Decoupling exploration and exploitation for meta-reinforcemmen learning without sacrifices. In International conference machine learning, pages 6925 6935. PMLR, 2021.
[28] D. Barber and Agakov. The im algorithm: a variational approach to information maximiza- tion. Advances in neural information processing systems, 16(320):201, 2004.
[29] R. E. Wang, M. Everett, and J. P. How. R-maddpg for partially observable environments and limited communication. arXiv preprint arXiv:2002.06684 2020.
[30] Yang, A. Nakhaei, D. Isele, Fujimura, and H. Zha. Cm3: Cooperative multi-goal multi- stage multi-agent reinforcement learning. arXiv preprint arXiv:1809.05788. 2018.
[31] J. Foerster, G. Farquhar, T. Afouras, N. Nardelli, and S. Whiteson. Counterfactual multi-agent policy gradients. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
[32] R. Lowe, Y. I. Wu, Tamar, J. Harb, O. Pieter Abbeel, and I. Mordatch. Multi-agent actor- critic for mixed cooperative-competiine environments. Advances in neural information pro- cessing systems, 30, 2017.
10

---

[33] S. Li. Y. Wu, X. Cui, H. Dong, F. Fang, and S. Russell. Robust multi agent reinforcement learn- ing via minimax deep deterministic policy gradient. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 4213-4220 2019.
[34] Y. Xiao, X. Lyu, and C. Amato. Local advantage actor-critic for robust multi-agent deep rein- forcement learning. In 2021 International symposium on multi-robot and multi-agent systems (MRS), pages 155-163. IEEE, 2021.
[35] A. Xie, L. M. Bhamidipary, E. Z. Liu, Hong. S. Levine, and C. Finn. Learning explore in pomdps with informational rewards. In Forty-first International Conference Machine Learning.
[36] S. C. Ong, S. W. Png, D. Hsu, and W. S. Lee. Planning under uncertainty for robotic tasks with mixed observability, The International Journal of Robotics Research, 29(8):1053-1068, 2010.
[37] H. H. Nguyen, A. Baisero, D. Klee, D. Wang, R. Platt, and C. Amato. Equivariant reinforce- ment learning under partial observability, In Conference on Robot Learning, pages 3309-3320. PMLR, 2023.
[38] D. Wang, C. Kohler, X. Zhu, M. Jia, and R. Platt. Bulletarm: An open- robotic ma- nipulation benchmark and learning framework. In The International Symposium of Robotics Research, 2022.
[39] D. Wang, R. Walters, X. Zhu, and R. Platt. Equivariant Q learning in spatial action spaces. In Conference on Robot Learning, pages 1713-1723, PMLR, 2022.
[40] H. Van Hasselt, A. Guez, and D. Silver. Deep reinforcement learning with double g-larnning In Proceedings of the AAAI conference artificial intelligence, volume 30, 2016.
[41] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, pages 1861-1870. PMLR, 2018.
[42] T. Ni, B. Eysenbach, F. SeyedSalehi, M. Ma, C. Gehring, A. Mahajan, and P.-.. Bacon. Bridg- ing state and history representations: Understanding self- seff-predictive rl. In The Twelfth Interna- tional Conference on Learning Representations, 2023.
[43] K. Cho, B. Van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Ben- gio. Learning phrase representations using rnn encoder-decoder for statistical machine trans- lation. arXiv preprint arXiv: 1078, 2014.
[44] M. Hausknecht and P. Stone. Deep recurrent Q-learning for partially observable MDPs. In 2015 aaai fall symposium series, 2015.
[45] V. Mnih, P. Badia, M. Mirza, A. Graves, Lillicrap, T. Harley, D. Silver, and K. Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International Conference Machine Learning, pages 1928 1937. PMLR, 2016.
[46] K. Perlin. Improving noise. In Proceedings the ics and interactive techniques, pages 681-682, 2002.
the 29th annual conference on Computer graph-
[47] D. P. Kingma and M. Welling. Auto-encoding variational bayes. 131iv: 14, 2013.
arXiv preprint
[48] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint urXiv:1707.000 2017.
[49] D. Wang, R. Walters, and Platt. SO(2) equivariant reinforcement learning. In International Conference Learning Representations (ICLR), 2022.

---

A Details of Agents
13
A.I
DTQN, TSAC
13
A.2
ZP-DRQN, ZP-RSAC
13
A.3
BA-DTQN, BA-TSAC
13
A.4
UA-DTQN, UA-TSAC
14
A.5
B-DQN, B-SAC
14
A.6 Hyper-parameters
16
A.7 Network Structures
17
Details of Domains
18
B.1 Sphinx
18
B.2 CarFlag-2D.
18
B.3 Heaven-Hell.
19
B.4 Robot Domains
19
Representations Training Details
C.I Training Data Generation
21
C.2 Network Architecture
21
C.3 Mutual Information Estimation
C.3.1 Minimizing [(=;;²²)
21
C.3.2 Maximizing [(;; = and 7 2x;
23
C.4
Hyper-parameters
23
D Additional Experiments
24
D.I Using x" @ d versus 2" for Task Learning.
24
D.2 Using Only Auxiliary Task/Intrinsic Rewards
D.3 Using GRU v.s. GPT
25
D.4
Visualization of Intrinsic Rewards
25
Details of Hardware Experiments
25
E.1 Obtaining Depth Images
25
E.2 Added Perlin Noise for Better Sim-To-Real Transfers
26
Details of SO(2) Rotational Data Augmentation
26
12

---

A Details of Agents
A.1 DTQN, TSAC
These are variants of DQN and SAC, made memory-based by using transformer as the sequence model shown in Fig. 8 and Fig. 9. Similar models have been explored in previous work [21. 22].
Concal
Qe aa)
0t
Encoder
Transformer
Q (14, an)
History-based Critic
Encoder
DTQN
0
Encoder
Figure 8: Architecture of DTQN.
Encoder
Transtormer
History-based Actor
Encoder
0
Encoder
TSAC
MLP
Qt (hr, a))
Encoder
Transformer
Q (he,
History-based Critic
Encoder
0g
Oe
Encoder
Figure 9: Architecture of TSAC.
A.2 ZP-DRQN, ZP-RSAC
These agents [42] are similar to DTQN and TSAC, except they use a recurrent sequence model in- stead of a transformer. Importantly, using a recurrent sequence model (e.g.. GRU [43]) is required (see [42]). Additionally, these agents regularized with seff-predictive auxiliary task of predict- ing the next latent state z from history h. Specifically, given recurrent encoder fø H 2 and a latent dynamics model ga Z x A + 2, the auxiliary task is to minimize:
Lasx (f) la
(6)
where T is the target network of Φ.
A.3 BA-DTON, BA-TSAC
In BA-DTQN [6] (see Fig. 10), state-based critic Q(s, a) and history-based critic Q(h, a) are leamned to leverage the state availability during training but not during exccution (i.e., state is available therefore cannot use Q(s,a) during execution). Unfortunately, pointed out in [6],
13

---

(1e ou)
Encoder
MLPP
State- State-based Critic
(s),
BA-DTQN
Cancal
Qec mu)
0
Enccdder
Transtomer
MLP
(h, oz)
History-based Critic
0g
Encodder
Encoder
Figure 10: Architecture of BA-DTQN.
Concal
0:
Encoder
Transtomer
MLP
History-based Actor
Enoocier
BA-TSAC
MLP
Q (4r, ax)
Encoder
B
State-based
Critic
07 (f 4)
4
Encoder
Figure 11: Architecture BA-TSAC.
Q(s,a) a) is not mathematically well-defined and is generally biased estimate of Q(h, a), which is used to select actions during execution.
The difference between BA-TSAC (see Fig. 11) and TSAC is that the critic is trained additionally using state input during training. Specifically, we learn a state-based critic Q(s, a) instead of history-based (h,, a). Similar BA-DTQN, BA-TSAC also has bias. For BA-TSAC, during execution, actions are computed using a history-based actor.
A.4 UA-DTQN, UA-TSAC
In UA-DTQN [6] (see Fig. 12), a history-state-based critic Q(h, s, a) and a history-based critic Qhh, a) are learned. Unlike BA-DTON with Q(s, a). Q(s, h, a) can be well-defined and has been proven be unbiased estimate of Q(h, a). During execution, actions are selected using Q(h, a).
Unlike BA-TSAC, UA-TSAC [9] (see Fig. 13) combines both state and history features train the critic, i.c., we learn history-state-basee critic Qs,, a). Similar to UA-DTQN, UA-TSAC does not introduce learning bias. During execution, actions are computed from history- policy.
A.5 B-DQN, B-SAC
The architectures of these agents are depicted in Fig. 14. These agents are based on Believer [8], which leveraged the state availability to train an agent in three stages:
Stage 1. Learning compact state representations with state-labeled transitions, i.e., a batch of sam- ples (s, o, a, r, s' o'). This stage is similar our first stage (see Algorithm but without the information-base regularizations. Instead, the authors proposed to regularize the KL divergence [L] 1)] to avoid overlapping features between d(s) and (b(o) by giving penalty when-
14

---

Encoder
Q: sp)
Encodier
Transformer
MLP
History-Slate-basee
(Re, ac)0))
Critic
Encoder
Encoder
UA-DTQN
Concat
Q (k, az)
Encoder
Transformer
MLP
History-based Critic
Encoder
Encodor
Figure 12: Architecture of UA-DTQN
Concal
Eacodee
Transtorm
History-based Actor
Ry-
Encoder
UA-TSAC
MLP
@h m.)
0g
Enooder
Transtrmer
MLP
(b, 4:0))
Enooder
Eistory-State-besse Critic
9:
Encoder
Encoder
Encoder
Figure 13: Architecture of UA-TSAC.
ever (bs) is used to derive features. This, however, does not avoid the overlapping issue between learned state and observation features, as shown in experiment (see Fig. 4).
Stage 2. Learn recurrent history model p(o(s)(h) with variational autoencoders [47] (VAE) by
maximizing the joint log-likelihood p(((()) h)) averaged over (s, h) samples.
Use the history module p(p(s) [h) for task learning. First, samples are drawn from the VAE derive a history summary. Then, this summary is used as the "states" for task learning using memoryless RL algorithms. The authors optionally fine-tune p(o())]k) with the on-policy data.
Stage 3.
As the original paper applied their method for PPO [48], which is on-policy, we had to modify the method to apply to DQN and SAC, resulting in B-DQN and B-SAC. In Stage 2, to fairly compare with other baselines, we replace GRU in the history model with the GPT model used in other base- lines. Moreover, in Stage 3, fine-tuned the history module for every domain (as used in the original code). Finally, the sequence model of B-SAC is shared between the actor and the critic, following the original code.
15

---

Pretrained History Encader
Concal
Qh,,, oe)
Transformer
Q14,, ax)
History-based
Critic
B-DQN
Op
Encoder
Pretrained History Encoder
MIP
Q (h, me)
History-based
Concat
Actor
Transformer
MLP
History-based
B-SAC
Critic
Encoder
Encoder
Figure 14: Architectures of B-DQN and B-SAC with a pre-trained history encoder from Believer [8]. We change the history encoder from GRU-based to transformer-basee for a fair comparison with other agents. For B-SAC, we shared history module, similar to the original code.
A.6 Hyper-parameters
For DDQN [40], we use an epsilon-greedy exploration strategy with a linear schedule, starting at e = 1.0 and ending at e = with T' being the episode length. The schedule time is equal to 10% of the total training timesteps. We use batch size of 64 episodes. For continuous actions, we use SAC [41]. We automatically tune the entropy temperature, initializing at 0.01. The chosen target entropy is equal the negation of the action dimension. We the discount factor y 0.99. We batch size of 64 episodes. Other hyper-parameters are in Table 2 with shared parameters and ones specific for each agent.
Table 2: Hyper-paramette parameters used for RL agents, HH: Heaven-Hell, S: Sphinx, CF: CarFlag-21
Agent
Hyper-parameter
Value
Shared among all agents
Episode Length
50
Discount Factor
0.99
Replay Buffer Size
1M: grid world domains, 100k: robot domains
Target Update Rate
0.005
Actor Learning Rate
3e-4
Critic Learning Rate
3e-5: CF and 3e-4: other
Batch size
64
SAC
Initial Entropy Temperature Update Per Step
0.01
0.25: discrete domains, 1.0: robot domains
ZP-DRQN
Loss weighting
1.0: discrete domains, 0.1: robot domains
Ours
Loss Weighting
0.5 for all domains
Reward Weighting
10.0: 0.1: robot domains, 0.0: CF
B-DON, B-SAC
Latent Dimension
32
X-Dim
16
Z-Dim
16

---

A.7 Network Structures
FC(n):: a fully connected layer with n outputs; Conv(ʃ, s): a convolutional layer with filter size f x and stride s; R: the ReLU activation function; MaxPool(w):: a max pooling layer with window size w: T(H, HS, D): Transformer with H heads, layers, hidden size HS, and the dropout rate D; GRU(N, HS): GRU with N layers and hidden size HS.
Branch
Encoder
=ogay
Action Encoder
FCB84)
Sequence Model
Tranit 1, 128, 6.3)
GRUIL 120)
Obs. Encoder
Pippa0g +1
(C)220) *M
Popeo
Figure 15: Network structures used Heaven-Hell.
Top Branch
Obs. Encoder
18 Comvt222+ +
22 ComvRa)+
Ratten
Encoder
Sequence Model
Obs. Encoder
64 Damegt13)-1 R
MasPwolt)
128 Curnett) #
Mastholgg
200 Donegt 33-11
Flatton
Figure 16: Network structures used in CarFlag-23 -D.
Top Branch
Obs, Encoder
18 Comp2.4)
Conv222)
Flarten
Action Encoder
Sequence Model
GRup 128)
Obs. Encoder 64 Comr2.0) +t
128 Come)2.3) +t
NosPool2)
256 Convi2.31
Flatonn
Bottom Branch
Figure 17: Network structures used in Sphinx.
Top Branch
Obs. Encoder
18 Comv(t11) R
52 Conv(23) +
Fi Flatten
FC)100)
Action Encoder
FC((4) R
Sequence Model
Transja, 128, 0.13
or
GRUIL 129)
Obs. Encoder
32 Conv/8.4) R
Conv(4.A) +
Convea.1)
Flatten
FC
Bottom Branch
Figure 18: Network structures used in robot domains.
17

---

B Details of Domains
B.1
Sphinx
-0.2
Channel
Charnel 1
Channel
Charnel 1
Observation
State
Figure 19: Sphinx domain with two-channeled pixel-based observations and states. Channel 1 of the observation reveals the goal cell (green) only when the agent enters the blue cell. In contrast, the same channel of the state always reveals the goal cell regardless of the agent's position.
In this domain (see Fig. 19), agent must visit the goal cell, which can be in one of three corners except the top-left one. The agent must visit the information cell (blue) the top-left corner to know the current corner of the goal. However, there is cost when going to the information cell.
Action. Move-Right, Move-Left, Move-Up, Move-Down
Observation. A 6 x 6 2 image with the first channel encodes the agent's position and the second encodes the goal's position. The second channel only contains the goal information when the agent enters the blue cell.
State. state has same structure as an observation, but the second channel always contains goal information.
Reward. +1 when reaching the goal, -0.2 when visiting the information cell, and otherwise.
B.2 CarFlag-2D
Channel
Channel
Channel
Channel
Observation
State
Figure 20: CarFlag-22 2D domain with two-channeled pixel-based observations and states. Channel of the observation reveals the goal cell only when the agent enters the blue region. contrast, the same channel of the state always reveals the goal cell.
In this domain (see Fig. 20), agent must visit the goal cell (green) finish the task. The goal cell, however, is only present in the observation when the agent visits the information region (blue).
Action. Right, Move-Left, Move-Up. Move-Down
Observation. A 11 11 image with the first channel encodes the agent's position, and the second encodes the goal's position. The second channel only contains the goal information when the agent enters the blue region.
State. state has the same structure as an observation, but the second channel always contains the goal information.
Reward. +1 when reaching the goal, and 0 otherwise.
18

---

B.3 Heaven-Hell
In this domain, an agent must visit heaven (green cell) to finish the task. The goal cell can be either on the left or on the right side with 50% probability. To observe the side of the goal (left or right), the agent must visit the priest, who resides in the bottom right corner.
-1
50%
50%
(s. y, (a, y true-side)
Agent
Priest
Figure 21: The SeaavenHHell domain with vector-based observations and states.
Action. Move-Riggee Right, Move-Left, Move-Up. Move-Down
Observation. vector consists of the agent's position and the side information. The side informa- tion can take the value of (no information), (heaven the right), -1 (heaven the left).
State. Like the observation, but the true side of the goal is always revealed.
Reward. +1 when reaching heaven, when reaching hell, and otherwise.
B.4 Robot Domains
In these domains, the agent must manipulate the only movable object among two objects, which exactly the same under the top-down depth image observation.
Action. An action a (6g, ky. 82, 6r), where days E [-0.05, 0.05] are the displacements of the gripper in the XYZ. axes, and 5, E (-m/8, x the angular rotation around the Z. axis.
Observation. All robot domains share the same observation: the top-down depth image taken from the camera centered at the gripper's position. Two fingers of the gripper are projected on the image.
State. The state also has two channels. The first channel the first top-down depth image of the observation. While the second channel of the observation non-informative, the second channel of the state an image that masks everything except the movable object (see Fig. 22) which the movable objects are colored red for visualizations).
Reward. In Block-Pulling, the agent receives a reward of 1.0 only when the two blocks are in contact. In Block-Pushing, the agent receives reward 1.0 only when the movable block is within em from the center the goal pad. Drawer-Opening, the agent receives a reward of 1.0 only when the unlocked drawer opened more than 5 em.
19

---

Channel 0
Channel
Channel
1
State
5
Channel
Channel 1
State
Figure 22: Visualization of an observation and a state in Block-Pulling, Block-Pushing, and Drawer-Opening. The movable object the red one. The state and the observation have two channels, the first being the top-down depth image. In Flock-Pulling the second channel of the reveals the movable object and the gripper. In Block-Pushing, the second channel reveals the movable block, the gripper, and the goal pad. In Brawer-Openingg the second channel in the state reveals the unlocked drawer and the gripper.
20

---

C Representations Training Details
C.1 Training Data Generation
Heaven-Hell, CarFlag-2D 2D, Sphinx: In these domains, we use uniform random agent generate training samples, each is transition (s,o,a,r,,'',") For the number of samples used in each domain, please see Table
In the robot domains, use the number of demonstrations (80 episodes) to learn the represen- tations during task learning. Furthermore, augment the training data using random rotations per transition as used in [49] (also see Appendix F). Finally, describe the planners used to generate the demonstrations in these domains.
Planner in Block-Pull ling: The planner randomly selects a block and attempts to pull it to the other block direction until the task is accomplished. If, for while, the position the selected block remains unchanged, the planner will move the gripper to the other block and repeat the pulling.
Planner in Block-Pushings A block is randomly chosen and pushed toward the goal pad. If the block's s position remains unchanged for while, the planner will move the gripper to the other block and resume pushing until the task finished.
Planner in Drawer-Openingg The planner selects drawer randomly and tries open it. If the chosen drawer fails to open after while, the gripper will move to the other drawer and repeat the opening action.
C.2 Network Architecture
The specific architecture used to learn representations shown in Fig. 23.
Maximize MI
Reward Predictor
State Encoder
f
Next Obs. Feature
Minimize
Sum
Dynamics Model
20
Predictor
Next State Feature
Obs. Encoder
Action Encoder
Predictor
Maximize MI
Figure 23: Architecture to learn representations in all domains.
Next, we describe the components for each domain from Fig. to Fig. 27. To succinctly describe the network architecture, we use the following acronyms: FC(n): fully connected layer with outputs; Conv(/, s): convolutional layer with filter size ʃ and stride s, R is the ReLU activation, and MaxPool(w): max pooling layer with window size w.
C.3 Mutual Information Estimation
C.3.1 Minimizing =²)
From the upper bound equation Eq. (2), minimize its variational estimate defined below:
1
LELUB
log '(>?!=f) log <<>g1=²))
B
21

---

State Encoder
Canv(2.1)
Conv(2.3)
MaxPool(2)
Conv(2,1)
FC(32)
Obs. Encoder
Action Encoder
FC(16)
Reward Pred.
Next State Feature Pred.
Dynamics Model
FC(258)
FC(258)
FC(256)
R
Next Obs. Feature Pred.
FC
Figure 24: Network architecture in Sphinx.
State Encoder
conug2.3) R
MaxPool(2)
120 Conv(211)
MaxPool(2)
Comv(211
Flatten
Obs. Encoder
Action Encoder
FC(16
Reward Pred.
FCD]
Next State Feature Pred.
FGG22)
Dynamics Model
FC]28) R
FC(26) +
FC(28))
Next Obs. Feature Pred.
FC(32)
Figure 25: Network architecture in CarFlag-2D.
State Encoder
FC(120) R
FC(120)
FC(32)
Obs. Encoder
Action Encoder
FC(16)
Reward Pred.
FC
Next State Feature Pred.
FC(32)
Dynamics Model
FC(128)
FC(128)
FC(128)
Next Obs. Feature Pred.
FC(22)
Figure 26: Network architecture
State Encoder
64 Conu(4,4)
Flatten
Obs. Encoder
Action Encoder
FC(16)
Reward Pred.
FG(1)
Next State Feature Pred.
FC(32)
Dynamics Model
FC(256)
FQ(256)
FC(256) +
Next Obs. Feature:Pred.
FC(32)
Figure 27: Network architecture robot domains.
22

---

The variational distribution (<="|2" is updated to minimize DxL 1g(>91=*) || p(=9|=*)] We assume (==1==)) follows Gaussian distribution and use the following network architectures:
Mean network: FC(32) R FC(32)
Log variance network: FC(32) - R + FC(32) Tanh
We use the batch size = 500 and use a learning rate of 0.001 for all tasks, except Heaven-Hell, which a learning rate of 0.0003 is used. We update q whenever we update ((s) and vb(0).
C.3.2 Maximizing 7 ((; and 7 (88; =*)
State Encoder
Observation Encoder
Layer
Layer
el Layer
x
Layer 0
Layer
Layer
olo,f) =
Layer
Layer
Layer1
Layer
4
Figure 28: Architecture calculate (s, E(s)) and a(o, E(o)) using the dot product operation.
From the lower bound equation Eq. (3), minimize the following loss:
B
CDIM
(p, sp(o((, E(x,))
1
As shown in Fig. 28, the discriminator a uses the same architecture of the state encoder & (when calculating the state feature zf of s) and the observation encoder /i (when calculating the observation feature zf o). We dot product to compute (s, E(s)) = of 2 and (o, E(o)) = 2f 2
C.4 Hyper-parameters
We provide the hyper-parameters used for training representations in Table 3.
Table 3: Hyper- Syper-parameters used in learning representation. HH: Heaven-Hell, S: Sphinx, CF:
CarFlag-20 2D, BP: Block-Pulling: BPs: Block- Plock-Pushing, and DD: Srraer-Openiine
Domain
BP
BPs
DO
# of samples
21785
45406
13682
1226
1240
1234
# of episodes
500
1000
500
80
80
80
# of augmentations per sample
4
12
6
# of training epochs
1000
1000
1000
1000
1000
1000
Batch size B
500
500
500
500
500
500
Learning rate
0.003
0.001
0.001
0.001
0.001
0.001
Reward loss coeff. A,
10.0
1.0
10.0
10.0
100.0
100.0
State loss coeff. A,
1.0
1.0
0.5
0.1
1.0
1.0
Observation loss coeff. X
0.5
5.0
0.03
1.0
1.0
1.0
I ((2):2") loss coeff. ACLUB
1.0
10.0
0.3
10.0
0.001
1.0
1 2(s;2*) loss coeff. ADIM
0.0
0.0
0.0
0.1
0.01
0.001
1 [(o 2)) loss coeff. ADIM
1.0
1.0
0.5
1.0
1.0
1.0
23

---

D Additional Experiments
D.1 Using 2 + 2" versus 2 for Task Learning
Continuing the experiment from Section 5.2.1, report the performance using 2 and 2 + 2 for task learning in all domains.
10
0.3
0.8
Recum
0.8
0.00
1V. 1
0.4
-0.25
0.1
-0.50
0.2
-2.75
666"
7000
0.0
100k
200w
3007
400K
500k
7038
4036
2002
Eivfronment
Stop
(a) Sphinx
(b) CarFlag-22 2D
(c) Geaaven-Helle
1.0
1.0
0.0
0.11
0.8
0.6
Retum
0.6
12.4
0.4
0.2
0.2
0.2
0.0
60k
Enwirowment Step
Environment Stop
Environment Stop
(d) Block-Pulling
Pusking (c)
Drawer-Opening
Figure 29: Task learning performance when using 2" and 2 @ 2 as the "state".
D.2 Using Only Auxiliary Task/Intrinsie Rewards
Here, we show the learning performance when using intrinsic rewards and/or the auxiliary task.
1.25
0.75
0.8
0.00
Bctum
0.6
Intrintt Only
Only
0.50
0.4
-025
0.2
-0.25
AuaBary-Dr
And Onl
-0.25
Acal Eon
0.0
-0.50
Aud woudayyeinninn
10oK
2005
300k
400K
500k
200k
400k
EOOK
800k
1.00
100%
2006
3000
400K
500K
Envtronment Stop
Environment Stop
Environment Stop
(a) Sphinx
(b) CarFlag-2D
(c) Heaven Hell
1.0
Vestral--Oll
3.0
Audi kary Only Fuod bey-ittiineee
0.8
3.8
0.6
2.6
0.4
0.4
0.2
pncrec Only Audibary-Clly
0.2
wmirdd-Olly Audilary-nnty
0.0
(d) Block-Pulling
(e) Block-Pushing
(f) Prraver-Opening
Figure 30: Comparing using intrinsic rewards the auxiliary task versus using both.
24

---

D.3
Using GRU v.s. GPT
Here, report the performance in all domains when using GRU versus GPT as the sequence model in our proposed agent.
GFT
0.5
0.6
10 b
My
0.4
0.2
GPT
0.0
CH0
1002
2002
2002
4002
5006
2002
4033
C00R
B000
2.0M
103%
2008
2002
400k
5006
Ensfronment Step
Environment Stop
Environment Stop
(a) Sphinx
(b) CarFlag-2D
(c) Heaven-Hell
1.0
DFU
CPT
CRJ
1.0
CHU
0.5
0.6
0.6
0.4
0.2
0.0
0.0
0.0
Ensironment
Environment Step
Enstronment
(d) Block-Pulling
(c) Pucki-n
(f) Draver-Opening
Figure 31: Task learning performance when using GRU v.s. GPT.
D.4 Visualization of Intrinsic Rewards
We visualize the intrinsic rewards of trained agents in three grid-world domains in Fig. 32. The intrinsic rewards peak when the agents perform the information-gathen actions.
0.30
Eptsode
Heaven
0.25
Henver Right
Ack Pricst
Intridic Reward
Enter Rogion
0.20
0.4
0.15
Epirade
0.2
0.10
Epicode
Epikade
0.0
0.05
-2
Enter Coll
0.00
Enstronment
Episode
Episode
(a) Sphinx
(b) CarFlag-22 2D
(c) dleaven-Hell
Figure 32: Intrinsic rewards within an episode of trained agents in three grid-world domains. Red
circles denote when the intrinsic rewards peak, c.g., when they perform informative actions.
E Details of Hardware Experiments
E.1
Obtaining Depth Images
We fuse the point clouds from two RealSense D455 cameras (Cam and Cam 2) and one Azure Kinect camera (Cam to create an integrated point cloud (see Fig. 33). We then orthographically project the point cloud at the gripper's position create depth image observation. Examples of observations in the three robot domains can be seen in Fig. 34.
25

---

Cam
Point
2
Cam
Orthographic Projection
Figure 33: We fuse the point clouds from three cameras (to avoid occlusions) and performed an orthographic projection at the gripper's position to create depth image observation.
(a) Block-Pulling
(b) Block-Pushing
(c) Srawer-Opening
Figure 34: Examples of observations in real robot experiments.
E.2 Added Perlin Noise for Better Sim-To-Real Transfers
Following [39], we found it useful for better sim-to-real transfers by adding the Perlin [46] noise to the depth images during training for more robust policies by being closer to real-world depth images. For all robot domains, we applied the noise with magnitude of 7mm (see Fig. 35).
Block-Pulling before.
Block-Pushing before.
Drawer-Opening before.
Block-Pulling after.
Block-Pushing after.
Drawer- Drawer-Qpening after.
Figure 35: Depth images before and after adding Perlin [46] noise for better sim-to-real transfers.
F Details of SO(2) Rotational Data Augmentation
We perform SO(2) rotational augmentation by choosing random angle and rotating the depth im- ages around its center. We perform this augmentation in two cases:
When learning the representations to utilize the data better. For each transition (s, o, a, r, s' o')), sample random angle and rotate 8, 0, s', d' at the same angle. Each transition has its random angle, see Fig. 36 for examples.
26

---

When performing task learning robot domains, Given an episode, we first sample random angle and apply the rotation with this angle for every 8, o, 8', of within the episode. Because we are trying to learn a history-basee policy, this is to ensure the augmented history is valid (see Fig. 37).
t=1
t= 2
t=3
1=
Figure 36: Examples of rotation data augmentation applied for transitions in an episode in Block-Pushing to augment the data for learning the representation: different random rotation is applied independently for s o' in each timestep in an episode.
1=1
1=
t
t=
t=3
Figure 37: Examples of rotation data augmentation applied for an episode in Block-Pushing: the random rotation consistently applied to every tuple s, o, s', of within episode.
27


Reviews:
[{"review_id": "THzudfGG54", "reviewer": "Reviewer_p5c4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "Work provides a MI-based Objective to encourage exploration in partially observable situations.", "summary_of_recommendation": "Provide a novel approach that outperforms baselines presented. Some issues with clarity and potential lack of scalabilty.", "technical_quality": 3, "clarity_of_presentation": 2, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "- One limitation seems to be that the first stage needs high quality demonstations. It would be interesting to consider how an imitation learning baseline (e.g diffusion policy) does with these demos as a point of comparison. This could also be further fine-tuned with the reinforcement learning using some ablation over the intrinsic/extrinsic reward to compare with the representation learned in P1. To deal with the partial observability using some history (either observations/actions) may be needed.\n- The discussion about the two branch architecture and Figure 1 are slightly unclear to understand how the concatenation is done and how it encourages the separation of $z^s$ and $z^0$. I understand that the objectives (e.g minimizing MI) can lead to this property but it isn't clear to me why the architecture choice can lead to this. Could you please clarify this?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Klc0K7tba7", "reviewer": "Reviewer_TSYf", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "This paper introduces an intrinsic reward bonus that encourages an RL agent to perform active information-gathering in partially-observable settings. The main idea is to assume access to the hidden states during training and introduce an exploration bonus to the policy for gathering information about these hidden states. The method is demonstrated on grid world problems and robotic manipulation tasks.", "summary_of_recommendation": "Overall, the paper presents a new method for learning exploration behaviors in partially-observable settings and demonstrates that it can indeed learn such behaviors in specific scenarios that call for them. However, some main flaws of the paper is that it doesn’t properly attribute certain derivations to prior work, and the experimental domains and exploration strategies required are quite simple.", "technical_quality": 3, "clarity_of_presentation": 4, "potential_impact": 3, "robotics_focus": 2, "quality_of_limitations": 3, "questions_for_rebuttal": "- How is the hyperparameter $\\alpha_t$ chosen? Why is it a function of the episode time-step?\n- What happens if the privileged state is incorrectly specified or is still a partial observation?\n- How would meta-RL methods compare on these experimental domains?\n- What are the main failure modes in the sim-to-real transfer experiment?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aAiluaDTzu", "reviewer": "Reviewer_VPtp", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "The paper proposes a novel method for learning from asymetric information in partially observable", "summary_of_recommendation": "The paper presents a convincing idea, with some minor issues and details in the experimental setup that seem easy to clarify in the rebuttal.", "technical_quality": 3, "clarity_of_presentation": 4, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 2, "questions_for_rebuttal": "As a simple additional baseline, I would ask the authors to run their method _without_ a history (or only a very short history) to highlight how much benefit comes from being able to aggregate information across multiple timesteps and how much is simply due to the architectural improvements.\n\nI would ask the authors to repeat the sim to real experiments with the strongest baselines as well to have a better basis for comparison.\n\nI would like to see the issues discussed in the limitations section to be highlighted and discussed in more depth earlier in the paper, especially how it influences the choice of evaluation environment and whether some baselines might outperform the proposed method on suboptimal data.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Pre-rebuttal:
This paper received 3 fairly confident reviews. The paper tackles learning under PO; they propose using state-observation and state-history mutual information; they use privileged training mechanisms and perform an ablation study for the different design decisions considered.

To address in rebuttal:
- Discuss sample complexity
- Address the nature of the chosen tasks (they seem contrived)
- Address limitations of imitation learning approach
- Improve on clarity (see reviewer comments)
- Better situate the work (eg, within Meta learning)
- Discuss the robustness of the current sim-to-real transfer strategy

Post-rebuttal:
The authors have addressed the issues pointed out by reviewers, and provided an improved manuscript; in particular, they have added experiments that further clarify the contributions of the method, juxtapose it more clearly with state of the art baselines, as well as the sim-to-real transfer robustness.
Decision: Accept