Leveraging Mutual Information for Asymmetric Learning under Partial Observability

Authors: hai  nguyen
long  dinh
robert  platt
christopher  amato
Venue: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-06 04:45:35

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.

Main paper: 
 introduction: 
 While partial observability is ubiquitous in robotics [1, 2] , most reinforcement learning (RL) robotics research avoids it due to the difficulty of training a policy that can effectively memorize past information and perform information-gathering actions. Fortunately, full or near-full observability during training is possible in many real-world domains using a simulator or a generative environment 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-based critic [6] in value-based RL to extract actions from. In contrast, using an information-theoric approach, we instead use the state to generate intrinsic rewards and optimize an auxiliary task, shaping the agent's behavior directly.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 a useful intrinsic reward to encourage informative actions. Moreover, we train the agent to predict the state features from past histories as an 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 supervised 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-based sequence model to output a history summary with maximal mutual information with the learned state features 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 experiments 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. 
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.A POMDP is defined by a tuple (S, A, Ω, T, R, O), where S, A, and Ω are the state space, the action space, and the observation space, respectively. In a POMDP, the transition function T (s, a, s ′ ) governs state changes, and the observation function O(a, s ′ , o) 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 action-observation history h t = (o 0 , a 0 , . . . , a t-1 , o t ) [12] .Mutual Information (MI) measures the dependence between two random variables x and y:EQUATIONwhere p(x, y) is the joint distribution and p(x), p(y) are the marginals.Minimizing MI: Cheng et al. [13] proposed to minimize I(x; y) by minimizing an upper bound named Contrastive Log-ratio Upper Bound (CLUB), which is formulated as:EQUATIONIn practice, p(y|x) is often unknown and approximated by a learnable variational distribution q(y|x).Maximizing MI: Instead of maximizing I(x, y), Deep InfoMax (DIM) [14] is a method that maximizes the MI between x and a high-level representation z x of x. The process involves an encoder E : X → Z x and a classifier σ : X × Z x → R to discriminate between samples from p(x, z x ) and the product of marginals p(x)p(z x ). DIM uses the Jensen-Shannon MI estimator [15] :EQUATIONwhere sp(z) = log(1 + e z ) is the softplus function. Maximizing I(x, z x ) involves continuously updating E and σ by maximizing the estimator I DIM through gradient ascent.In practice, we optimize Eq. ( 2) and Eq. (3) through samples; see the Appendix for more details.We focus on the asymmetric RL setting [3, 9, 6, 8] in which, during training, both the state s t and the observation o t are available while only o t is accessible during execution. Unlike an MDP, the goal here is to find a history-based (not reactive) policy π(a t |h t ) which maximizes the expected discounted returnJ = E [ ∞ t=0 γ t R(s t , a t )],where γ ∈ [0, 1) is a discounting factor. We use an information-based approach to leverage states during training for efficiently learning π(a t |h t ). 
method: 
 Our method has two stages. First, we learn non-overlapping state and observation features from a fixed dataset containing privileged state information. The dataset is generated using a random agent or a simple hard-coded planner. Next, we use the learned features to train a memory-based agent with a two-branch architecture using information-based intrinsic rewards and an auxiliary task.This stage aims to learn state features z s (encoded by a state encoder ϕ) and observation features z o (encoded by an observation encoder ψ) with two desired properties. (P1): Firstly, we want the concatenated feature z s ⊕ z o to encode task-relevant state features. In other words, z s ⊕ z o is a compact "state", which is useful when the raw state is high-dimensional or contains task-irrelevant features. (P2): Secondly, we want z s to contain unobservable task-relevant features, while z o contains observable task-relevant features. With P1 and P2 satisfied, z s will encode only task-relevant but unobservable features inferred from the history as z o already contains features extracted from the observation. We can also view z s as the outcome of an ideal information-gathering process, which results in task-relevant hidden information not captured by the current observation (i.e., z o ).We satisfy (P1) by training z s ⊕ z o given an action a to be predictive of z s ′ , z o ′ and the reward r, similar to bi-simulation studies for MDPs [16, 17, 18] . For (P2), Wang et al. [8] minimized the KL divergence between ϕ(s) and a fixed standard normal distribution by giving a penalty whenever features are derived from the s 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 I(z s ; z o ) using Eq. ( 2). However, minimizing only I(z s ; z o ) may result in an undesired scenario, where I(z s ; z o ) is small but only ϕ(s) or ψ(o) captures all good features, and the remaining contains no useful information. Therefore, we further regularize learning with maximizing I(o; z o ) and I(s; z s ) using Eq. (3). Algorithm 1 summarizes this stage (more details in the Appendix C).Algorithm 1 Learning z s and z o (more details in Appendix C)Require: Dataset D = {(s i , o i , a i , r i , s ′ i , o ′ i )} N i=1 1: repeat 2:Sample a batch of data{(s i , o i , a i , r i , s ′ i , o ′ i )} B i=1 ∼ D 3: Compute current features z s i = ϕ(s i ), z o i = ψ(o i ) 4:Compute reward and next features using a dynamics model g: ri , ẑis ′ , ẑi o ′ = g(z s i , z o i , a i ) 5:Compute target featuresz s ′ i = ϕ(s ′ i ), z o ′ i = ψ(o ′ i ) 6: Compute L r = (r i -ri ) 2 , L s = ∥stop-grad(z s ′ i ) -ẑi s ′ ∥ 2 2 , L o = ∥stop-grad(z o ′ i ) -ẑo ′ i ∥ 2 2 7:Compute L CLUB using Eq. (2) to minimize I(z s ; z o ) 8:Compute L DIM using Eq. (3) to maximize I(o; z o ) and I(s; z s ) 9:Update ϕ, ψ, g to minimize λ r L r + λ s L s + λ o L o + λ CLUB L CLUB + λ DIM L DIM 10: until convergenceAfter learning z s and z o , we now introduce a 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 z s , which are not captured in z o . For the bottom branch, we directly use ψ(o) to initialize it. A two-branch architecture is not 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 z s and z o and using z s to bias the top branch's learning, we alleviate the top branch's memory burden by training it to focus only on features not readily observed.Because z s is the ideal outcome of the information-gathering process, we encourage informationgathering by rewarding actions that produce h t that can maximize I(h t ; z s t ). Moreover, we use maximizing I(h t ; z s t ) as an auxiliary task to help memorize important information from history. Encouraging Information Seeking. By maximizing I(h t ; z s t ), the agent is encouraged to gather information by acting to get h t that encodes maximal information with z s t . We closely follow [27] and maximize the mutual information by maximizing the following lower bound [28] :EQUATIONwhere f θ (z s |h) serves as an approximation to the true conditional distribution p(z s |h). Similar to [27] , we assume f θ (z s |h) follows a Gaussian distribution centered around a deterministic history embedding f θ (h) with variance ρ 2 I. We parameterize f θ (h) using the GPT-based sequence model and train f θ (h) to predict z s across batches of (h, z s ) samples by minimizing:EQUATIONSince only the third term of Eq. ( 4) depends on π, we interpret logf θ (z s t |ht) f θ (z s t |ht-1)as an intrinsic reward term r intrinsic t . This intrinsic reward captures per-step incremental information gain that the current history h t reveals about z s t compared to the previous h t-1 . With a weighting factor α, our final reward therefore becomes r total t := r env t + αr intrinsic t . In practice, we calculate the intrinsic rewards asr intrinsic t = d cosine (z s t , f θ (h t ))-d cosine (z s t , f θ (h t-1)), with d cosine being the cosine similarity distance. Encouraging Information Retention. While the intrinsic reward encourages information gathering, we also want the agent to memorize important observations (or actions). For this purpose, we train the agent's top branch (k in Fig. 1 ) to predict z s over a batch of episodes as an 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(z s |h)).Training the agent this way can be considered an indirect and "online" way to maximize I(h; z s ) so that the agent can mentally memorize important events within the history h that increased I(h; z s ). 
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, 32, 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 observability, Belief-grounded 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 a history-based critic. Using the actor-critic framework, Baisero and Amato [9] learn a history-state-based critic to help learn a history-based actor. Concurrent with our work is [35] , which also proposes informational rewards to encourage information gathering. While these methods use the privileged information as additional input (e.g., to train a critic), we first find a compact and task-relevant representation of the privileged information (similar to [8] ). However, we diverge from [8] by directly using privileged information for encouraging information gathering. Our information-based approach shares similarities with a 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 a single policy using a combined reward and add an auxiliary loss to improve memorization. Our idea of minimizing the overlapping between z s and z o has connections with the mixed-observability framework [36] , which removed the overlapping between state and observation to reduce the computation for state distribution 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).In Sphinx [8] , CarFlag-2D [37] , and Heaven-Hell [6] (see Fig. 2 ), an agent ( ) must reach a randomly generated goal position ( ). Specifically, the goal can be one of the green corners in Sphinx, left/right ends in Heaven-Hell, and any cell in CarFlag-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 the blue cell in Heaven-Hell, the side of the goal (left/right) is observed. Similarly, in the Sphinx and CarFlag-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 ( ) if any, and 0 otherwise. In Sphinx, an additional penalty of -0.2 is given when entering the blue cell, making the information-gathering process costly. In these domains, the state encodes both the current observation and the goal cell.In Block-Pulling, Block-Pushing, and Drawer-Opening (modified from domains in [38] ), a robot must pull, push, or open the only movable object (a block or a drawer) among two, which appear the same under a 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-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 to pre-populate the replay buffers of all agents. Moreover, because these tasks exhibit spatial symmetry [39] , we use random SO(2) rotations (see Appendix F), applied consistently for all observations and actions within an episode, for data augmentation.We use DDQN [40] (discrete action spaces) and SAC [41] (continuous action spaces) as the base RL 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 a transformer (a simplified 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 a transformer 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, Q(s, a) and Q(h, a) are learned during training, and argmax a 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 a history-based actor. These asymmetric agents can sometimes perform well [3] but were found [6, 9] to introduce learning bias, i.e., Q(s, 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, s, a) (instead of Q(s, a)) and Q(h, a) and use Q(h, a) to pick actions during execution. B-DQN, B-SAC: Believer [8] originally trains a GRU-based history encoder, similar to our first stage. For a fair comparison, we use the same GPT to 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 RL algorithms, i.e., DQN or SAC in our case. Furthermore, we fine-tune the history encoders to strengthen these baselines, following the public code. 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, Block-Pulling, and Drawer-Opening. Most surprisingly, Believerbased baselines (B-DQN/-SAC) do not perform well in all domains despite our significant effort to strengthen them (replacing GRU with GPT, hyper-parameter tuning, and online fine-tuning). When information seeking is hard (e.g., in Sphinx with a clarifying penalty, or in Heaven-Hell with Figure 4 : a) Overlapping test: using z s or concatenating z s and z o for task learning; b-c) Learning curves in Heaven-Hell when using our method (Section 3.1) or Believer [8] to learn z s and z o . a 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-2D, Block-Pulling, and Drawer-Opening) thanks to pre-trained compact features.To show that our method can achieve both (P1) and (P2) in Section 3.1, we perform task learning in Heaven-Hell using a) z s ⊕ z o , b) only z s , and c) raw state s. We compare two ways to learn the representation: our method and Believer [8] . From Fig. 4 , z s ⊕ z o yields better performance than the raw state. This indicates that our method can learn high-quality state features. As expected, using z s alone fails to solve the task because z s will only contain the features of the goal's position. The task can only be solved with the aid of z o , which contains the agent's position. In contrast, when using Believer [8] to learn z s and z o , using either z s ⊕ z o or z s makes no difference: both can solve the task, indicating that z s also contains the agent's position, which is already contained in z o .In this section, we perform ablation studies to analyze our key design choices: a) using GPT instead of a recurrent neural network (e.g., a 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 Appendix D for more details. Using GPT instead of a recurrent neural network. While a GRU [43] might have trouble memorizing important information in long sequences (i.e., saw important observations but forgot), a GPT can attend to and memorize any past information. This property is crucial for our agent to maximize I(h; ϕ(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 retention, 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 is quite large) and the robot domains (agents are provided with demonstration episodes), using 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. To verify the performance of the learned policy in simulation, we perform sim-to-real transfers using the best policies in simulation to a UR5 robot (see Fig. 7 ). To best obtain non-occluded topdown 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 1 show that the learned policies can be zero-shot transferred in the real world (see our supplementary video for more details), and our 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 Drawer-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: 
 Conclusion. This work proposed an information-based method that directly uses the mutual information between the state and history to encourage information-gathering under partial observability. We also tackle memorization by incorporating an auxiliary task to predict the state feature from the history to alleviate the agent's memory burden by only requiring its memory component to memorize 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.Limitations. Our method's biggest limitation is the challenge of learning good and non-overlapping representations at the first stage. In practice, we need to select the right dataset to train with. We can use random interactions in the grid-world domains, but expert demonstrations are needed for training in the robot domains so that the representations are useful. In practice, to guarantee the learned representations can satisfy both (P1) and (P2), we need to iterate between hyper-parameter 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. Q(s, a) is not mathematically well-defined and is generally a 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 the history-based Q(h, a). Similar to BA-DTQN, BA-TSAC also has bias. For BA-TSAC, during execution, actions are computed using a history-based actor. 
 
 Appendix: In UA-DTQN [6] (see Fig. 12 ), a history-state-based critic Q(h, s, a) and a history-based critic Q(h, a) are learned. Unlike BA-DTQN with Q(s, a), Q(s, h, a) can be well-defined and has been proven to be an 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 to train the critic, i.e., we learn a history-state-based critic Q(s, h, a). Similar to UA-DTQN, UA-TSAC does not introduce learning bias. During execution, actions are computed from a history-based policy.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 samples (s, o, a, r, s ′ , o ′ ). This stage is similar to our first stage (see Algorithm 1) but without the information-based regularizations. Instead, the authors proposed to regularize the KL divergence KL[ϕ(s)∥N (0, 1)] to avoid overlapping features between ϕ(s) and ψ(o) by giving penalty when- ever ϕ(s) is used to derive features. This, however, does not avoid the overlapping issue between learned state and observation features, as shown in our experiment (see Fig. 4 ).Stage 2. Learn a recurrent history model p(ϕ(s)|h) with variational autoencoders [47] (VAE) by maximizing the joint log-likelihood p(ϕ(s), h)) averaged over (s, h) samples.Stage 3. Use the history module p(ϕ(s)|h) for task learning. First, samples are drawn from the VAE to 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(ϕ(s)|h) with the on-policy data.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 baselines. Moreover, in Stage 3, we 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. We change the history encoder from GRU-based to transformer-based for a fair comparison with other agents. For B-SAC, we use a shared history module, similar to the original code.For DDQN [40] , we use an epsilon-greedy exploration strategy with a linear schedule, starting at ϵ = 1.0 and ending at ϵ = 1 T with T being the episode length. The schedule time is equal to 10% of the total training timesteps. We use a 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 to the negation of the action dimension. We use the discount factor γ = 0.99. We use a batch size of 64 episodes. Other hyper-parameters are in Table 2 with shared parameters and ones specific for each agent. In this domain (see Fig. 19 ), an 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) at the top-left corner to know the current corner of the goal. However, there is a cost when going to the information cell.Action. Move-Right, Move-Left, Move-Up, Move-Down Observation. A 6 × 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. A state has the same structure as an observation, but the second channel always contains the goal information.Reward. +1 when reaching the goal, -0.2 when visiting the information cell, and 0 otherwise. In this domain (see Fig. 20 ), an agent must visit the goal cell (green) to finish the task. The goal cell, however, is only present in the observation when the agent visits the information region (blue).Action. Move-Right, Move-Left, Move-Up, Move-Down Observation. A 11 × 11 × 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 region.A 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.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. State. Like the observation, but the true side of the goal is always revealed.Reward. +1 when reaching heaven, -1 when reaching hell, and 0 otherwise.In these domains, the agent must manipulate the only movable object among two objects, which are exactly the same under the top-down depth image observation.Action. An action a = (δ x , δ y , δ z , δ r ), where δ xyz ∈ [-0.05, 0.05] are the displacements of the gripper in the XYZ axes, and δ r ∈ [-π/8, π/8] is 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 is the first top-down depth image of the observation. While the second channel of the observation is non-informative, the second channel of the state is an image that masks everything except the movable object (see Fig. 22 ) in 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 a reward of 1.0 only when the movable block is within 5 cm from the center of the goal pad. In Drawer-Opening, the agent receives a reward of 1.0 only when the unlocked drawer is opened more than 5 cm. Here, we show the learning performance when using intrinsic rewards and/or the auxiliary task.  


Decision: 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.

Reviews: review: title: title: Review for Submission 256
paper_summary: paper_summary: Work provides a MI-based Objective to encourage exploration in partially observable situations.
main_review: main_review: Strengths:
- The authors present a novel MI objective grounded by prior work.
- The authors show efficacy of the method across domains and show a substantial improvement over alternative approaches
- The authors show the transfer of their approach to real domains while training in sim.
- The authors perform a rigorous ablation study for the different design decisions considered

Weaknesses:
- Would be good to consider the sample complexity of an algorithm, given the first stage requiring a high quality dataset to train with as stated in the limitation section.
- The tasks considered seem slightly contrived and would be interesting to see how this approach works for more complicated domains
- The approach has many moving components which draws up concerns about scalability and efficacy when the data, architecture, and algorithm is scaled up with the objective functions designed.
- Can improve clarity of paper (e.g some work in figures 1,4) could be helpful for readers to better understand the approach.
questions: questions: - 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.
- 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?
limitations: limitations: 3
review_summary: review_summary: Provide a novel approach that outperforms baselines presented. Some issues with clarity and potential lack of scalabilty.
title: title: Review of Leveraging Mutual Information for Asymmetric Learning under Partial Observability
paper_summary: paper_summary: 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.
main_review: main_review: Strengths:
- This paper introduces a novel exploration bonus that encourages information-seeking behaviors, by assuming access to ground-truth states during training.
- The method is demonstrated on a series of simulated tasks that require exactly such information-gathering behaviors, and it outperforms prior methods that do not have such exploration explicitly incorporated into them.
- Paper is written clearly.

Weaknesses:
- I don't think the paper is appropriately citing what is new in this paper and what has been introduced in prior work already. Specifically, the section on "Encouraging Information Seeking", which discusses the derivation of the reward bonus, is taken nearly directly from Liu et al [1]. I think it should be attributed accordingly. Additionally, I wanted to point out recent/concurrent work that's probably of interest to the authors. It introduces a similar exploration bonus for the same problem setting [2].
- Access to privileged state information even during training can be an overly strong and restrictive assumption.
- Due to this assumption, the method is probably most useful for scenarios where we can first train in simulation, which more readily offers such privileged information, and transfer to real.
- The current experimental domains require pretty simple information-gathering strategies, and could be casted as meta-RL problems. How would meta-RL methods like [1], [3], [4] compare here?

[1] Liu et al. Decoupling exploration and exploitation for meta-reinforcement learning without sacrifices. ICML 2021.

[2] Xie et al. Learning to Explore in POMDPs with Informational Rewards. ICML 2024.

[3] Zintgraf et al. Varibad: A very good method for bayes-adaptive deep RL via meta-learning. 2019.

[4] Kamienny et al. Learning adaptive exploration strategies in dynamic environments through informed policy regularization. 2020.
questions: questions: - How is the hyperparameter $\alpha_t$ chosen? Why is it a function of the episode time-step?
- What happens if the privileged state is incorrectly specified or is still a partial observation?
- How would meta-RL methods compare on these experimental domains?
- What are the main failure modes in the sim-to-real transfer experiment?
limitations: limitations: 3
review_summary: review_summary: 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.
title: title: Official review
paper_summary: paper_summary: The paper proposes a novel method for learning from asymetric information in partially observable
main_review: main_review: Quality: The paper is overall of a high quality. The writing and figures are well done, and the theory (albeit limited) is correct and well explained.

Clarity: The paper is cleanly written with minimal jargon and no unnecessary formalism. The extensive appendix is useful to get information on the experimental setup which is slightly terse in the main paper.

Originality and Significance: The use of privileged information for policy learning is relatively well established in the literature. However, the paper makes a meaningful contribution by showing improvements across a wide range of baselines on some tasks. I think a full evaluation of the potential impact is difficult as the set of tasks is somewhat limited and the performance improvements are not fully clear

## Open issues

Overall, the idea seems to be well supported in the experiments. However, I believe that the ablations are incomplete and deserve more attention. The introduced method changes several details compared to the compared baselines and the source of the improvements are unclear.

The ablation clearly shows that a majority of the performance of the method seems to stem from the use of a GPT style transformer as opposed to the GRU cells, with the intrinsic reward also contributing some performance. However, i.e. the DTQN baseline is stated to use DQN while the proposed method uses DDQN. Since this is an orthogonal issue, I would ask the authors to rerun ablations with a more clear and equal footing. Furthermore, ZP-DRQN etc are claimed to require a GRU, but I am not convinced that this is true. While Ni et al state that transformers do not fulfill the recurrency condition, I am not convinced that the test environments require a full recurrent state encoder. The baseline comparison/ablation requires substantial care here, and the quoted paper talks about a more general set of problems (i.e. POMDPs with in principle unbounded history memory requirements) which the use of transformers in this method would be unable to solve.

I want to highlight that the authors were otherwise very careful with the baseline comparison, so this is a bit nipicky.

This problem also applies to the intrinsic reward. I would like to ask the authors to comment on which baseline methods could be augmented with the same reward. I think it is important to not label arbitrary combinations of orthogonal advancements as their own "algorithms" in reinforcement learning, but to embrace the "rainbow lesson" and evaluate heuristics and improvements as individual improvements which can be freely combined.

The robotics experiments are somewhat short, given the focus of the conference, I am unsure whether the method is best placed at CoRL. However, since the authors do include both some arguments for why the problem is relevant for robotics and some real world experiments, I am not weighing this concern strongly.
questions: questions: 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.

I would ask the authors to repeat the sim to real experiments with the strongest baselines as well to have a better basis for comparison.

I 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.
limitations: limitations: 2
review_summary: review_summary: The paper presents a convincing idea, with some minor issues and details in the experimental setup that seem easy to clarify in the rebuttal.