Learning Task Belief Similarity with Latent Dynamics for Meta-Reinforcement Learning
Authors: Menglong Zhang
Fuyuan Qian
Quanying Liu
Venue: ICLR 2025 Poster (2025)
Keywords: meta-reinforcement learning
representation learning
bisimulation
Primary area: reinforcement learning
TL;DR: We propose SimBelief—a novel meta-RL framework via measuring similarity of task belief in Bayes-Adaptive MDP (BAMDP).

Abstract: Meta-reinforcement learning requires utilizing prior task distribution information obtained during exploration to rapidly adapt to unknown tasks. The efficiency of an agent's exploration hinges on accurately identifying the current task. Recent Bayes-Adaptive Deep RL approaches often rely on reconstructing the environment's reward signal, which is challenging in sparse reward settings, leading to suboptimal exploitation. Inspired by bisimulation metrics, which robustly extracts behavioral similarity in continuous MDPs, we propose SimBelief—a novel meta-RL framework via measuring similarity of task belief in Bayes-Adaptive MDP (BAMDP). SimBelief effectively extracts common features of similar task distributions, enabling efficient task identification and exploration in sparse reward environments. We introduce latent task belief metric to learn the common structure of similar tasks and incorporate it into the real task belief. By learning the latent dynamics across task distributions, we connect shared latent task belief features with specific task features, facilitating rapid task identification and adaptation. Our method outperforms state-of-the-art baselines on sparse reward MuJoCo and panda-gym tasks.

Published as a conference paper at ICLR 2025
LEARNING TASK BELIEF SIMILARITY WITH LATENT DYNAMICS FOR META-REINFORCEMENT LEARNING
Menglong Zhang, Fuyuan Qian, Quanying Liu" Southern University Science and Technology (zhangm 12022, qi ananfy22233 @ma wwwtuch.euu.......
1 Liugy@ssstech.ee edu.cc.
ABSTRACT
Meta-reiiforcemmet learning requires utilizing prior task distribution information obtained during exploration to rapidly adapt to unknown tasks. The efficiency of an agent's exploration hinges on accurately identifying the current task. Recent Bayes-Adaptive Deep RL approaches often rely reconstructing the environ- ment's reward signal, which is challenging sparse reward settings, leading to suboptimal exploitation. Inspired by bisimulation metries, which robustly extracts behavioral similarity in continuous MDPs, we propose SimBelief a novel meta- RL framework via measuring similarity of task belief in Bayes-Adaptive MDP (BAMDP). SimBelief effectively extracts common features of similar task dis- tributions, enabling efficient task identification and exploration in sparse reward environments. We introduce latent task belief metric to learn the common struc- of similar tasks and incorporate into the specific task belief. By learning the latent dynamics across task distributions, we connect shared latent task be- lief features with specific task features, facilitating rapid task identification and adaptation. Our method outperforms state-of-the-ar baselines on sparse reward MuJoCo and panda-gym tasks.
INTRODUCTION
In meta-reinfforemmett Icarning, an agent is required to efficiently explore the environment gather information relevant the current task and use that information to adapt to new, unseen tasks (Duan al. 2016; Finn al. 2017 Gupta et al.] 22018; [Humplik al.] 2019 However, in real-world scenarios, there are many distrac- tions, and irrelevant information can cause the agent to engage in erroncous exploration behaviors, leading to the learning of a sub- optimal policy. This challenge exacerbated when the exploration space is vast, and rewards are sparse. Humans. on the other hand, have the ability to generalize across similar tasks by quickly identi- fying common patterns, such recognizing that both opening win- dow and a drawer involve a pulling action. These shared structures can be abstracted and applied to new tasks through certain transfor- mations, enabling rapid adaptation by leveraging prior knowledge. This process involves identifying similar features and relationships between tasks, with the underlying task belief facilitating effective knowledge transfer (Figure T
Latent Space
MDP
MDP
MDP k
Real Environment
Taski
Task
Taski
Learning shared structures among tasks the latent space as task beliefs enables the agent to rapidly adapt to new tasks.
Figure I:
Common approaches to task belief modeling include posterior sam- pling (Thompson) 1933; Osband et al., 2013 Rakelly et al. 2019)
and Bayes-Adaptive Markov Decision Processes (BAMDPs) (Duff 2002 Ghavamzadeh et al.) 2015), where BAMDPs provide a structured framework for learning and adapting to new tasks by integrating prior information (Zintgraf et al. 2019). Previous works have addressed the issue of
"Corresponding author.

---

Published as a conference paper at ICLR 2025
uninformative rewards in meta-RL by using intrinsic rewards (Zhang et al. (021c) extract task- relevant information or exploration bonuses (Zintgraf et al.) 2021 to enhance performance in hard exploration tasks. However, they often neglect the potential common structures shared across tasks.
To efficiently learn the structural relationships embedded within task belicf and address the chal- lenges of effective task identification and rapid adaptation, it is necessary develop methods that can efficiently recognize tasks with similar structures. This requires exploring strategies be- yond BAMDP- approaches, which may involve leveraging hierarchical reinforcement learning (Frans et al. 2018), or task embeddings (Liu et al. 2021 Lan et al. 2019) that can capture the underlying structure of tasks and facilitate transfer learning across related tasks.
In this work, we propose SimBelief' (Learning Similarity of Task Belief for Meta-RL) as meta-RL framework under BAMDP to measure the difference of tasks in similar distributions by learning the dynamics and transferring capability in latent space. This results in a task belief similarity, which quantifies the representational relationships between any two random tasks in the learned latent space. However, directly reconstructing and exploring the environment in the latent space may lose crucial detailed information (Yarats et al.] 2021; Kemertas Aumentado- Artastoon..r [021] Hafner et al. 2019), which is important for accurately reconstructing the specific task and ensuring convergence to the Bayesian optimal policy(Zintgrnn et al. 2019 Choshen & Tamar 2023).
Therefore, combine the task belief similarity learned in the latent space, which provides an overall understanding of the task distribution and the relationships between tasks, with the belief of the specific task currently being learned, significantly enhancing the agent's is adaptability unknown environments. Compared to related meta-RL methods, our approach effectively learns dynamics in the latent space and leverages the learned latent representations to quickly discern relationships between tasks. This enables the agent to utilize prior knowledge from tasks, recognize the current environment, and accelerate exploration and adaptation unseen tasks.
Our main contributions are as follows:
We propose task representation method for context-based meta-reinforcement learning al- gorithms in BAMDP, which enhances tasks by learning task belief similarity through latent reward model, transition model, and inverse dynamics model.
the agent's ability recognize and adapt unknown
In scenarios with sparse or uninformative reward signals, our algorithm more effectively extracts latent task space representations compared to other baselines and achieves state- of-the-art performance on sparse reward MuJoCo and panda-gym tasks.
Our algorithm demonstrates stronger adaptability and generalization capabilities out-of- distribution (OOD) tasks.
We theoretically validate the effectiveness of the latent task melief Markov Decision Processes (BAMDPs).
metric in Bayes-Adaptive
BACKGROUND
2.1 PROBLEM FORMULATION
In meta-reinforcement learning (meta-RL), consider distribution p(7) over a space of tasks T. where each task is modeled as partially observable Markov decision process (POMDP; (Cassandra et al. 1994)) defined as T (S, A, O, P. R. Po. 7, H). Here, S and A represent the shared state and action spaces across all tasks, while O denotes the task-specific observation space. The transition function P(s' 8, a), reward function R(r 8, a), and initial state distribution Po define the task- specific dynamics.
For context-based methods, the agent interacts with the enviromment over a horizon of H time steps in each episode of each task and utilizes a task-specific context the hidden task dynamics. The objective to learn policy T(a Tι) that can rapidly adapt to new tasks by leveraging the learned task belief encoded in the context. The full horizon comprises N task episodes, referred to as a meta- episode, each with its own set of T time steps, and the task context is updated as the agent gathers more observations. The objective in meta-RI. is maximize
(O).t, 01:t-1, (1a-1) to infer
2

---

Published as a conference paper at ICLR 2025
the expected return:
((x) = TTp(T
To address the multi-task adaptation problem in POMDPs, it is essential avoid local optima and balance exploration with exploitation. Utilizing the framework of Bayes-Adaptive Markov Decision Processes (BAMDPs; (Dufi) ).02), we can represent uncertainty over task-relevant information in a belief space. By leveraging the current belief, we can reconstruct the dynamics model to better handle uncertainty, enabling more effective adaptation to the current task. The BAMDP is defined as M = (S+ .A Pū P Rt H*), where the hyper-state space S = S x B combines the state space S with the task belief space B, representing the posterior over MDPs based on past experiences. p+ is the transition function, PT is the initial hyper-state distribution, and R is the reward function. The total decision horizon H+ N x И. The goal is to learn a meta-policy r+t(a, s that maximizes cumulative reward by balaneing exploration of task uncertainty and exploitation of the current belief.
2.2 BISIMULATION METRICS IN RL
Bisimulation metrics provide way to measure state similarity in reinforcement learning (RL) based on behavioral equivalence, focusing on transitions and rewards rather than raw state features. The concept originated in model checking (Givan et al. 2003). In RL, bisimulation aggregates states that lead to similar future outcomes. This allows for more efficient learning. Approximate bisimulation metrics (Ferns et al. 2004 Castro 2020) extended this idea to high-dimensional environments, allowing RL agents to better generalize, reduce state space complexity, and improve exploration. These metrics are particularly useful in tasks with sparse rewards or partial observability (Guo al., 2022).
Definition 1 (z-Bisimulation Metric (Castro 2020 Given two states s; and sj from the state space S, the distance between these states can be evaluated using the π r-bisimulation metric dx (St, 8z), which measures behavioral similarity. We define the metric d (818 85) as:
d 8j) R, RR,,1 + Wi (d.))PE PT,)
where Ry, and Re,, represent the reward functions under policy, pa denotes the transition distribu- functionss tions, and W1 the Wasserstein distance between the distributions.
2.3 UNINFORMATIVE REWARDS IN META REINRFRRRIEEEET LEARNING
Uninformative rewards provide insufficient feedback for agents to quickly learn effective strategies (Andrychowicz et al. 2017 Dulac-Arnole et al., 2019). This issue can impede exploration, slow convergence, and lead to suboptimal policies. To address this challenge, methods such intrinsic motivation and curiosity- spaces when extrinsic rewards are minimal (Pathak al., 2017 Burda al. 2018). Reward shaping techniques modify the reward function to offer more informative feedback, thereby accel- erating learning (Ng al. [999). Balancing the exploration-exploitation trade-off is also crucial for efficient learning (Sutton) 2018). Enhancing meta-RL. agents' ability to learn from uninforma- tive rewards can significantly improve their performance across complex, real-world environments by incorporating advanced exploration strategies and reward augmentation methods (Gupta et al., 2018b).
exploration have been employed to encourage agents to explore state
In this work, we learn the dynamics of multiple tasks in a latent space and use task belief metric to measure the differences between similar tasks. By extracting similar structures within the latent space and representing these structures as task belief similarity, we map them to the current real and unknown environment to assist the agent exploration. This approach aims to achieve rapid adaptation to multiple tasks in environments with uninformative rewards.
3 METHOD
In this section, we provide a detailed introduction to SimBelief (Figure 2). an effective off-policy meta-RI. approach for online adaptation within the BAMDP framework. First, we propose latent
3

---

Published as a conference paper at ICLR 2025
task belief metric to quantify the relationships between different tasks in the latent space in section Then, we explain how task belief similarity is learned within the learned latent dynamics in section 3.2 We also describe the overall algorithmic process of SimBelief in detail in section 3.3 Finally, we theoretically prove the conditions under which the policy can transfer between tasks in the latent space in section 3.4
3.1
3.1 LATENT TASK BELIEF METRIC
Previous work (Zhagg et al. 2021b) learns an environment encoder to capture the relationships between Block MDPs (Du et al. 2019) for multi-task generalization. In (Zhang et al., 2021c). in- formation gain is used as intrinsic reward for exploration. However, in these approaches, the agent needs know the current task information (i.e., task ID) during training. In contrast, in context-based meta-RL, these approaches limit adaptability to out-of-distribution tasks, the agent must infer the task being executed from historical information and reason about the current task distribution during adaptation. Including task ID during training can partially constrain the agent's reasoning capabilities. In real-world scenarios, relying solely on the environment's s inherent reward signal manually designed intrinsic rewards often fails to effectively capture the common struc- ture of tasks (Zheng et al.) 2020). Therefore, learning the structural similarities between tasks and leveraging belief facilitate rapid transfer across similar tasks is more advantageous in complex environments, especially those with sparse reward signals. To measure the relationships between tasks and make it applicable within the BAMDP, we define the latent task belief metric.
Definition 2 (Latent Task Belief Metric). Given two latent task beliefs ba and bj. with corresponding samples 2; and 2j identified by their latent representations, let dz denote the metric that evaluates the distance between these two task beliefs. We define da (z) as follows:
dz zz) = Ie, (st, a.) R75 R55 ag) | + W2(da) (17 T,),,, (j))
+ -17 1
where of = (g((s), zi) is the augmented state corresponding to the task r2 R3 (8 as)) is the reward model, StT T (st ag) is the transition model, and a, [5 (t stt) is the inverse dynamics model.
Here, we map s latent compressed space as g(s), and learn the task dynamics in the compressed space. Latent dynamic space can be defined as S G(S) x Z. In Figure 2 represents the sam- ple from latent task belief b; in the latent space and captures the similarity between different tasks P All task distributions share common latent dynamic space, with different tasks distinguished by different zt, where zι contains the dynamic similarity information between any two tasks. We em- ploy the inverse dynamics model to predict as it not only helps retain information relevant to what the agent can control but also enhances the agent's 's reasoning capability (see Appendix G.2). Unlike in Definition [ here we use the dynamic model measure task belief similarity to enhance exploration efficiency in unknown environments, while preserving the agent's control information in sparse reward settings. This approach is more effective for identifying similar tasks.
3.2 LEARNING TASK BELIEF SIMILARITY IN LATENT SPACE
Our goal is to use the latent task belief metric to capture the similarity information between any two tasks, which is embedded within the task belief. This enables SimBelief to achieve rapid task reasoning and adaptation. From the replay buffer DD we randomly sample trajectories from two tasks (i.e., task and task j) in the task set M. The trajectory then fed into the context encoder 4o to obtain historical information h and h², which serve as prior knowledge for the tasks and are input into the latent space. The latent space primarily consists of three components: (1) state encoder ge that maps the original state space the latent state space S+, (2) a belief similarity learner d that maps h to the latent task belief b, and (3) latent dynamics model po conditioned on the learned latent task belief (Figure
's' represents the state the next time step.
For simplicity, in the latent task helief metrie, #i 2j and 2f. of are used interchangeably to represent the meaning.
4

---

Published as a conference paper at ICLR 2025
Context Encoder
Real Dynamics Pe
Replay Buffer D
Tia, a. P
h!
h
vaB
9
s
M b z!
Belief Similarity Learner
9e
by
bf
g ((s{)
zj
g(s))
Latent Dynamics
ar
Reward Model
Transition Model
Inverse Dynamic Model
Free
R
1
Distm
Figure 2: SimBelief architecture. Our framework consists of three components: learning latent task belief similarity, reconstructing the specific task, and the policy. In the compressed state space, sh1 learns the latent task belief bf through latent dynamics shared among all tasks in the distribution. In the real environment, we reconstruct the dynamics of the specific task using t to obtain the specific task belief by. The policy ı is conditioned
the state and the combined belief b', which integrates bj. and bf. In this paper, we use I and r to denote the latent space and real space, respectively, and i and represent task i and task respectively.
To obtain the latent dynamic representation of all tasks conditioned on 2t, reconstruct the latent dynamics po((s'=, r.o) =), which includes the reward model R2, the transition model T,, and the inverse dynamics model I through one-step predictions. The learning objective for the latent dynamics can be expressed as:
By(M) (B)A(a) [logPo((s"* ,r, =)]]]
(4)
This objective enables the model to capture task-specific latent dynamics by maximizing the likeli- hood of predicting the next augmented state s reward r, and action given the latent task belief Z)-
We aim to learn the similarity relationship between any two tasks using the latent task belief metric and the latent dynamics learned through Equation 4 This similarity relationship obtained by optimizing the belief similarity learner 4/2. The learning objective for task belief similarity using yhy defined as follows:
Cbisim (wh) (hv) - (h1 - R(5T a4)) R(555 W2 (T(- s.a)), (- ((115,a;))
.855 I1,
(5)
where the gradients of dynamics model are stopped, and bj = (h;). We use the 2- Wasserstein metric because the W₂ metric has convenient closed-form solution (Zhang al. (02 During training, through Equation and (5) zi, sampled from bj. is optimized to encode both the latent dynamic information required reconstruct task and the similarity information between task and any other task in the task set M. Through this approach, we can utilize 8iz to identify contextual information the current timestep and represent the similarity between any tasks through its output br.
3.3 SIMBELIEF ALGURITHM
We now proceed to describe how the learned task belief similarity (Section 3.2 can leveraged to enhance the exploration capability and convergence efficiency of meta-RL. The algorithm pscu- docode can be found in Appendix C
Reconstructing the specific task. In the BAMDP framework (Ghavamzadeh et al. 2015 Zintgral et al. 2019), even with extremely sparse rewards, it is necessary to construct the specific task's reward model in the original state space using variational auto-encoder (VAE, (Kingma 2013)
This because reconstructing the specific task in the latent state space may result the loss of crucial information, as the latent space is solely used for learning task belief similarity. Additionally,
5

---

Published as a conference paper at ICLR 2025
along the temporal dimension, only reconstruct the past history, approach in the latent space. This approach effectively captures local information about the task distribution, enhancing the agent's reasoning ability for future predictions (see the upper part of Figure 2. The objective is to maximise
one-step prediction
[vAB(Ø) = (x.((e)me)) Pog -+ - DKL (go(2r Ta) p(+)))
(6)
where 96 generates the specific task belief 2, for the current timestep using only the information from the historical trajectory ти. We consider qo as the forward reasoning process to infer the task belief 2,, corresponding 2, (b, tt)g (h T.E) in the Algorithm
Integrating latent task belief for effective online exploration. The learned latent task belief 2 ~ b, enables the agent quickly recall the similarities between different tasks during interactions with the environment, enhancing the efficiency of reasoning and exploration. This is the primary reason why SimBelief achieves robust adaptation OOD tasks (Figure 4.. Since the latent state space is an abstraction of the original task, the latent task belief may overlook the specific task information required by the agent during exploration in the real environment. Therefore, we choose to project the two distributions into joint space B, and then use a Gaussian mixture of the two distributions 6,- (Hy, (h), u. (h)) and by (Hv, (h), ow (h)), combined with the real environment's augmented state derived from as the input to the policy æ. The formulation of the Gaussian mixture distribution b given by:
(7) where N( Hr, of) and N( of, Gaussian distributions with means μ4, and f4, and variances of and of respectively. The weights w, and wi satisfy w. wr = 1 and are non-negative (see Appendix G.3
b = wy - -2 Hr. (r: + wi (2 pr oi,)
In the early stages of training, the latent task belief enhances the agent's ability to identify tasks and explore efficiently in the real environment, providing high-level understanding of the overall task distribution (see Appendix (F2). However, for the algorithm to converge stably, it need to incorporate finer-grained information about specific tasks, which is captured in the specific task belief b,.. Therefore, we minimize the discrepancy between the latent belief by and the specific task belief b, when training thy. This is achieved by introducing a KL divergence term between the two beliefs, serving regularization for the latent task belief. Combined with Equation 5 the overall optimization objective of 4/z can be written as:
((tr) = Ep(ar) [Crisim (vh) + DRI. (Go(2, T.)) p(x)))
(8)
where the gradient of 2s is stopped, and p(z)) represents the distribution of 21 (r h).
SimBelief utilizes Soft Actor-Critic (SAC) (Haarnoja et al. 2018) to optimize the policy, where the optimization of latent dynamics and latent task similarity conducted jointly with policy opti- mization to ensure convergence to near Bayes-optimal policy. To enhance the consistency between latent dynamics and real environment exploration, without introducing perturbations that could dis- rupt belief similarity in the latent space, we use the loss of Q O-function as the objective to train the offset (Aμr, Ao) of b; (M (h) + Ap, v. (h) Ao) (see Appendix G.4b. The loss of SAC after integrating latent dynamic information is
(TT) E (s,a)> [a log (I(a (s) - (o b),
where (s, b) is the augmented state in the real environment under the BAMDP framework.
3.4 THFORETICAL ANALYSIS
Based the latent task belief metric d. (2i, 22) defined in Definition 2 establish the follow- ing theorems regarding the difference in task dynamics between two tasks in the latent space within BAMDP.
Theorem 1 (Value difference bound). Given two tasks M, and M2 in the latent space with states st. 8 E St, and let v be the value function of policy n, the value difference bound between the tasks can be given by:
(=(+t) 5d dz(,,
(10)
where d (Zi, (i) is the latent task belief metric.

---

Published as a conference paper at ICLR 2025
Sparee
Sparse
Walkas-Rand--Paard
200
200
Fransca
Frames
105
Frames
1et
Pande-Reach
Pande-
Parda -iac---Pic---PPPeee Aakk-Aad--iacee
1.0
0.5
0.6
0.8
04
8.8
Su00206 Rah
0.6
88
04
0.1
6.0
0.
Frumos
SimBelief (Cung)
PEARL
MetaCURE
VoriBAD
Hyperx
RL2
Figure 3: Meta-testing performance on sparse reward MuJoCo and panda-gym tasks over 3 random seeds. Our algorithm, SimBelief, demonstrates superior online adaptation capabilities compared to other algorithms.
This theorem demonstrates that tasks with similar latent task beliefs =; and will have similar value function under the same policy
Theorem 2 (Latent transfer bound). Let Q'M; be the optimal Q-function task Mj. The difference between Q',, and the Q-function of the policy n learned from task Mi applied to task Mj, is bounded as follows:
[Qm,- 1QMiMMl SER + (ET ter || -=5l11)
Rmax
(11)
2(1 71)
This extended transfer bound theorem provides formal bound how well policy trained on task M, will transfer to task Mj. ER, ET, and are approximation errors for rewards, transitions, and inverse dynamics, All proofs are provided in Appendix B
4 EXPERIMENTS
In this section, we evaluate SimBelief sparse-reward tasks in MuJoCo (Finn et al. 2017 Rakelly et al, 2019) and the more challenging panda-gym (Gallouédec et al., 2021) environment. We aim to address the following questions: 1. Can SimBelief achieve fast online adaptation in sparse reward tasks? 2. Can SimBelief leverage learned latent belief similarity representations to enhance out-of- distribution generalization? 3. What the impact of latent task representations on rapid exploration? 4. How does the latent space correspond to the real environment?
Environments and baselines: We conducted experiments on six complex sparse reward tasks, in- cluding Point-Robot-Spars:: and Panda-Pick-And-Place (see Appendix [C). While MuJoCo tasks are commonly used by current meta-learning algorithms, panda-gyn gydda-gymmmo simulates weal-world robotic arm movements with extremely sparse rewards, making it a better benchmark for evaluating an algorithm's 's ability handle real- world tasks. We compared SimBelief against PEARL (Rakelly et al.] (019), which is based on posterior sampling, and MetaCURE (Zhang et al. 2021e which learns a separate policy, as well as VariBAD /Zintgraf et al. 2019), HyperX (Zingrraf et al. 2021), and RL (Duan et al. 2016). which are based on the BAMDP framework (see Appendix D)
-eeeaa--e Valker-Rand-Parant Panda-Reach, Panda-Push.
Online Adaptation Performance.
During the training phase, we performed meta-testing by calcu- lating the meta-episode average return and success rate across different tasks to evaluate the algo- rithm's 's online performance. As shown in Figure SimBelief consistently performed well across all tasks and exhibited superior adaptation capabilities compared to other algorithms. BAMDP- algorithms often require more interactions with the environment converge Bayes-optimal

---

Published as a conference paper at ICLR 2025
Velociy Range: [0.0. 3.0g
Veladly Range: [3.0. 4.0)
Velocity Range: [4.0. 5.0)
350
2540
Avarage Hestuu
Avarage Raturn
120
200
120
30
Episodes
Episodos
Eploodos
Radius
Radius
Rodius
Nargag Ruoum
Epicodes
Estoodes
SimBelief (Ours)
PEARL
MetaCURE
VariBAD
HyperX
RL2
Figure 4: Average test performance for the first 5 rollouts on out-of-distribution tasks (Cheetah- Vel-Sparse and Point-Robot-Sy -Sprsse) shows that as the task distribution gradually deviates from the training distribution, SimBelief demonstrates more efficient and robust adaptation compared to other algorithms, while still maintaining a high return.
policy, whereas algorithm, by integrating task belief similarity in the latent space, significantly improved the agent's ability to effectively extract task-relevant information in sparse reward envi- ronments, thus accelerating convergence. In sparse reward tasks like those in panda-gym, the per- formance of Varibad, HyperX, and RI.² illustrates that overly relying on the environment's 's reward signal while neglecting task similarity information makes difficult to achieve online adaptation in more complex scenarios. Additionally, SimBelief demonstrated more stable performance compared to posterior sampling-bais based algorithms like PEARL and MetaCURE, indicating that SimBelief ac- curately learned the common structure of the task distribution.
Out-of-distribution Task Inference and Adaptation. We evaluated the algorithm's 's reasoning and generalization capabilities on -if---f-ddstibution (OOD) tasks. In Figure we show the agent's s out- of-domain adaptation and generalization performance on the Cheetah- Sppraaee (top) and Point- Robot-Sparse (bottom) tasks. The leftmost column represents the task range used during the training phase, where the agent was trained by generating task goals within this range. During the testing phase, adjusted the velocity range for Cheetah- Vel-a---eeeee-eeeeeeeee and the semi-circle radius for Point- Robot-Sparse to assess the agent's robustness and task inference capabilities within five episodes. The results indicate that SimBelief exhibits stable nut-of-distributio generalization, demonstrating that learned latent belief similarity representations can be leveraged to enhance out-of-distribution task inference ability. The latent task belief similarity proposed also equips the agent with strong task transferability.
Exploration Performance. We visualized the exploration and adaptation process of the agent in the Woint-Robot-Spatee environment over two episodes, as shown in Figure [ The test target ap- peared within a semicirele with radius of 1.2, and five positions on the semicircle were randomly selected as the agent's targets. Compared to other algorithms, SimBelief successfully identified the randomly appearing target location within the semicircle and controlled the robot to reach the tar- get point within two episodes. Moreover, in the more efficient learning of a near Bayes-optimal policy compared to BAMDP- algorithms like HyperX. The latent task belief enabled the agent efficiently analogize and identify tasks, leading to more effective exploration.
setting, SimBelief demonstrated
Task Belief Visualization. We use t-SNE (Van der Maaten & Hinton 2008) to visualize SimBelief's latent task belief and specific task belief (Figure aim to illustrate how the task common structure leamned in the latent space contributes to effective reconstruction of the real environment. This vi- sualization reveals the underlying reason for SimBelief's s ability to achieve rapid out- dirdiitrinttion generalization. The latent space, as an abstraction and holistic representation of the task distri-
8

---

Published as a conference paper at ICLR 2025
Hyperx
Episode
Figure 5: Exploration and adaptation performance in the random out-of-distribution (radius 1.2) tasks of Point-Rohie SonttRooot-Spare. SimBelief is capable of robustly learning near Bayes-optimal policy.
Sen beiid specitic odle!
Sincaladd Luseers bele!
Hasx
Figure 6: t-SNE visualization of task beliefs learned by the algorithm 10 randomly sampled OOD tasks. The velocity range for Cheetah- See-Sae is set to 3-4, and the semicircle radius for Point-Robot-Spanse 1.1. SimBelief demonstrates more distinct representation in the latent space, enhancing the specific task belief to represent similar tasks in the real environment.
bution, enables efficient transfer essential task information. This, turn, allows the agent to leverage prior knowledge for identifying and reasoning quickly about unknown environments.
The mechanism of SimBelief for rapid adaptation. SimBelief, through the latent task belief, essentially learns the transfer relationships between knowledge across different tasks. In extremely sparse reward scenarios, once the agent succeeds in one task, this prior knowledge of success can be quickly propagated other similar tasks via the latent task belief, enabling the agent rapidly adapt to similar types of tasks. The latent task belief acts as a bridge, facilitating the generalization of past successes to new, related tasks, significantly reducing exploration time and computational resources. A more detailed discussion is provided in Appendix
RELATED WORKS
Task representation in reinforcement learning focuses on how agents encode and utilize task in- formation to enhance learning efficiency and generalization. Meta-reinforcemen learning enables agents to quickly adapt to new tasks by learning task priors; for instance, (Finn et al., (017) intro- duced Model-Agnostic Meta-Learning (MAML.) for rapid adaptation with minimal updates, while

---

Published as a conference paper at ICLR 2025
Rakelly al. 2019) proposed PEARL to learn latent task representations for fast probabilistic adap- tation. (Zingrraf et al. 2019: Dorfman et al. 2021 1 learn a latent variable model task distribution for efficient addeetvee RL, and (Gupta et al. 2018a) explored unsupervised meta-learning with- out explicit task labels. Yuan & Lu ) Choshen & Tamar (2023) apply contrastive learning to enhance encoder representation. (Lee et al. 2023) decomposes complex tasks into subtasks han- dle con-paammetri task variability. In multi RL, (Teh et al. 2017)) introduced Distral to learn shared policies across tasks with task- sspcciii adaptations. Contextual task representations have also been explored: (Sodhani et al., (021)) use metadata to learn interpretable representations, and (Hausman et al., 2018)) learned embedding spaces for transferable skills. These studies underscore the importance of efficient task representation to improve RL performance and generalization.
Exploration with task inference. Integrating exploration with task inference in reinforcement learning enables agents to learn efficient policies by understanding task structures. Information- theoretic approaches encourage exploration by maximizing the entropy of state visitation distribu- tions, promoting diverse behaviors (Liu & Abbeel 2021). (Raileanu & Rocktäschel] 2020) propose a type of intrinsic reward which encourages the agent to take actions that lead to sigmfiicntt changes in its learned state representation. (Wan et al., 2023) introduce DEIR, method that enhances ploration efficiency and robustness by using discriminative model-based episodic intrinsic rewards. Yang al. 2023) propose unsupervised skill discovery method through contrastive learning among behaviors. (Rana et al.] 2023) propose a low-level residual policy for skill adaptation en- abling downstream RL agents to adapt to unseen tasks Planning- haaning-bbseed methods enable agents to plan exploratory actions that maximize information gain by learning world models (Sekar et al., 2020
Xie al. 2020) leverage the idea of partial amortization for fast adaptation at test time.
Bisimulation for control. Bisimulation has been effectively applied in various control tasks, demonstrating its value in reducing computational burdens without sacrificing policy performance (Zhang et al., (021a).. (Zhang et al., (021b) propose method for learning state abstractions that generalize across tasks, improving sample efficiency and performance in multi-task and meta- reinforcement learning by leveraging shared dynamics in complex environments. (Gelada et al.) 2019) has leveraged bisimulation for efficient exploration by ensuring the agent treats bisimilar states similarly, reducing sample complexity. Blansen-Estrucde et al. 2022) propose form of state abstraction that captures functional equivariance for coal-conditioned RL. (Kemertas & Aumentado- Armstrong 2021) generalize value function approximation bounds for on-policy bisimulation met- ries to non-optimal policies. These studies highlight the increasing applicability of bisimulation in reinforcement learning and control tasks. They offer variety of solutions that reduce sample complexity, guide exploration, and enhance learning performance, particularly in goal-conditioned settings and beyond.
CONCLUSION
We present SimBelief, meta-RL framework aimed at improving the agent's ability to quickly adapt and generalize real-world sparse reward tasks. At the core SimBelief is the latent task belief metric, which learns the similarity between the dynamics of various tasks in a shared latent space and represents this similarity as task belief similarity. This helps the agent efficiently capture the common structure of the task distribution, facilitating the recognition and reasoning about unknown tasks. By combining the latent task belief with the specific task belief learned through interactions in the real environment, SimBelief demonstrates strong adaptation and generalization capabilities during exploration, particularly on out-of- of-istribbution tasks. It overcomes the slow convergence and inefficiency issues inherent in BAMDP. We believe that SimBelief will inspire future research on agent generalization, providing a powerful tool for enabling agents to more efficiently tackle the challenges and complexities of real-world environments.
ACKNOWLEDGMENTS
This work wasfunded in part by National Natural Science Foundation of China (62472206), Shen- zhen Excellent Youth Project X Shenzhen Science and Technology In- novation Committee (2022410129,
KCXF72200122111744000
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAI Picter Abbeel, and Wojeiech Zaremba. Hindsight experience re- play. Advances in neural information processing systems, 30, 2017.
Yuri Burda, Harrison Edwards. Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv: 1810. 2018.
Anthony Cassandra, Leslie Pack Kaelbling, and Michael L Littman. Acting optimally in partially observable stochastic domains. In Aaai, volume 94, pp. 1023-1028, 1994.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 10069-10076, 2020.
Era Choshen and Aviv Tamar. Contrabar: Contrastive adyattvee deep rl. In International Conference on Machine Learning, pp. 5005-6021 PMLR, 2023.
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning- identifiability challenges and effective data collection strategies. Advances in Neural Information Processing Systems, 34:4607-4618, 2021.
Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient rl with rich observations via latent state decoding. In International Conference on Machine Learning. pp. 1665-1674. PMLR, 2019.
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:1611.02773 2016.
Michael O' O'Gordon Duff. Optimal Learning: Computational procedures for Bayes-adaptive Markov decision processes. University of Massachusetts Amherst, 2002.
Gabriel Dulac-Amold, Daniel Mankowitz, and Todd Hester. Challenges of real-world reinforcement learning. arXiv preprint arXiv: 2019.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite markov decision processes. In UAI, volume pp. 162-169, 2004.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leaming for fast adaptation of deep networks. In International conference on machine learning, pp. 1126-113 1726-1135 1135. PMLR, 2017.
Kevin Frans, Jonathan Ho, Xi Chen, Pieter Abbeel, and John Schulman. Meta learning shared hierarchies. In International Conference on Learning Representations, 2018.
Quentin Gallouédec, Nicolas Cazin, Emmanuel Dellandréa, and Liming Chen. panda-gym: Open- source goal-conditioned environments for robotic learning. In 4th Robot Learning Workshop: Supervuuee and Lifelong Learning@ NewrIPS 2021, 2021.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International conference on machine learning, pp. 2170-2179. PMIR, 2019.
Mohammad Ghavamzadeh, Shie Mannor, Joelle Pineau, Aviv Tamar, al. Bayesian reinforcement learning: A survey. Foundations and Trends@ in Machine Learning, 8(5-6):359-435 2015.
Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization in markov decision processes. Artificial intelligence, 147(1-2):163-223, 2003.
Yijie Guo, Qiucheng Wu, and Honglak Lec. Learning action translator for meta reinforcement learning sparse-reward tasks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 6792-6800, 2022.
Abhishek Gupta, Benjamin Eysenbach, Chelsea Finn, and Sergey Levine. Unsupervised meta- learning for reinforcement learning. arXiv preprint arXiv:/80e 046404 2018a.

---

Published as a conference paper at ICLR 2025
Abhishek Gupta, Russell Mendonea, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Meta- reinforcement learning of structured exploration strategies. Advances neural information pro- cessing systems, 31, 2018b.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01660 2019.
Philippe
Amy Zhang, Ashvin Nair, Patrick Yin, and Sergey Levine. Bisimula- tion makes analogies in goal-conditioned reinforcement learning. In International Conference on Machine Learning. pp. 8407-8426. PMLR, 2022.
Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin Riedmiller. Learning an embedding space for transferable robot skills. International Conference on Learn- ing Representations, 2018.
Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv: 1rXxv:1905006424 2019.
Nan Jiang. Notes state abstractions, 2018.
Mete Kemertas and Tristan numentado-Armstroon Towards robust bisimulation metric learning. Advances in Neural Information Processing Systems, 34:4764-4777, 2021.
Diederik Kingma. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 2013.
Lin Lan, Zhenguo Li, Xiaohong Guan, and Pinghui Wang. Meta reinforeement learning with task embedding and shared policy. Proceedings of the 28th International Joint Conference on Artificial Intelligence, pp. 2794-2800. 2019.
Suyoung Lee, Myungsik Cho, and Youngchul Sung. Parameterizing non-parametric meta- reinforcement learning tasks via subtask decomposition. Advances in Neural Information Pro- cessing Systems, 36:43356-43383, 2023.
Evan Z Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn. Decoupling exploration and ex- ploitation for meta-reinforcemen learning without sacrifices. In International conference on ma- chine learning, pp. 6925-6935, PMLR, 2021.
Hao Liu and Pieter Abbeel. Behavior from the void: Unsupervised active pre-training. Advances in Neural Information Processing Systems, 34: 18:199188477 2021.
Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transformations:
Theory and application reward shaping. Ieml, volume 99, pp. 278 277,287, 1999.
Ian Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. Advances in Neural Information Processing Systems, 26, 2013.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driver exploration by self-supervised prediction. International conference on machine learning, pp. 2778-2787. PMLR, 2017.
Roberta Raileanu and Tim Rocktäschel. Ride: Rewarding impact-driven exploration for procedurally-geneeeeee environments. In International Conference on Learning Representationss 2020.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off ooolpo renforreinor learning via probabilistic context variables. In International conference on machine learning, pp. 5331-5340. PMLR, 2019.
12

---

Published as a conference paper at ICLR 2025
Krishan Rana, Ming Xu, Brendan Tidd, Michael Milford, and Niko Sünderhauf. Residual skill policies: Learning an adaptable skill-based action space for reinforcement learning for robotics. In Conference on Robot Learning, pp. 2095-2104. PMLR, 2023.
Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In International conference on machine learning, pp. 5583-8592. PMLR, 2020.
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International Conference on Machine Learning, pp. 9767-9779, PMLR, 2021.
Richard Sutton. Reinforcement learning: An introduction. A Bradford Book, 2018.
Yee Teh, Victor Bapst, Wojciech M Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. Advances in neural information processing systems, 30, 2017.
William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-294 1933.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Shanchuan Wan, Yujin Tang, Yingtao Tian, and Tomoyuki Kaneko. Deir: efficient and robust
ploration through seweriiinaiieemmdel-baaeee episodic intrinsic rewards. In Proceedings of the Thirty-Seoond International Joint Conference Artificial Intelligence, pp. 4289-4298 2023.
Kevin Xie, Homanga Bharadhwaj. Danijar Hafner, Animesh Garg, and Florian Shkurti. Latent skill planning for exploration and transfer. arXiv preprint arXiv:2011.1388 2020.
Rushuai Yang, Chenjia Bai, Hongyi Guo, Siyuan Li, Bin Zhao, Zhen Wang. Peng Liu, and Xuelong Li. Behavior contrastive learning for unsupervised skill discovery. International Conference on Machine Learning, pp. 39183-39204. PMLR, 2023.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improv- ing sample efficiency in model-free reinforcement learning from images. In Proceedings of the anai conference artificial intelligence, volume 35, pp. 10674-10681, 2021.
Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reinforcemenn learning via contrastive learning. In International Conference on Machine Learning, pp. 25747-25759.
PMLR, 2022.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learn- ing invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021a.
Amy Zhang, Shagun Sodhani, Khimya Khetarpal, and Joelle Pineau. Learning robust state abstrac- tions for tions, 2021b.
block mdps. In International Conference Learning Representa-
Jin Zhang, Jianhao Wang, Hao Hu, Tong Chen, Yingfeng Chen, Changjie Fan, and Chongjic Zhang Metacure: Meta reinforcement learning with empowerment-drivee exploration. In International Conference on Machine Learning, pp. 12600-12610, PMLR, 2021c
Zeyu Zheng, Junhyuk Oh, Matteo Hessel, Zhongwen Xu, Manuel Kroiss, Hado Van Hasselt, David Silver, and Satinder Singh. What can learned intrinsic rewards capture? International Confer- ence on Machine Learning, pp. 11436-11446 PMLR, 2020.
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta- learning. arXiv preprint 1ixiv 2019.
13

---

Published as a conference paper at ICLR 2025
Luisa M Zintgraf, Leo Feng, Cong Lu, Maximilian Igl. Kristian Hartikainen, Katja Hofmann, and Shimon Whiteson, Exploration in approximate hyper-state space for meta reinforcement learning. In International Conference on Machine Learning, pp. 12991-13001. PMLR, 2021.
14

---

Published as a conference paper at ICLR 2025
A THEORETICAL BACKGROUND AND ANALYSIS
In this section, we introduce the theoretical background and the rationale behind the design of the latent task belief metric.
Definition 3 (Bisimulation Relation (Givan et al, 2003)), Given an MDP M. an equivalence re- lation E S S x S is bisimulation relation if whenever (s,t) E. the following properties hold, where Se is the state space S partitioned into equivalence classes defined by E:
Ya E A, A, = R(t,a)
(12)
Va A,YC SF. P(s,a))) P(t, ((t),
(13)
s/Ee
Definition (Diffuse Metric), A diffuse metric measures the distance between two points by consid-
ering not only the shortesi path but also the distribution of multiple paths between the points. The diffuse metric satisfies the following properties:
1. Non-negativity: For any two points and y, d(x, v) 2 0. and d(x, v) = 0 if and only if y- Symmetry: For any two points x and y, d(x, y) = d(y, z)..
3. Triangle inequality: For any three points y, and z, d(z, S d(x,3) + d(y.z).
These properties ensure that the diffuse metric behaves as a valid distance measure across the space.
To prove that the metric (2i, 2j) in Defination 2 is valid, we show it satisfies three key properties: non-negativity, symmetry, and the triangle inequality.
|R] a.) R2 (sJ az)l is the abso- lute difference in rewards, which non-negative (TT (s, T7)) .a)) is the Wasserstein distance, which is non-negative. [1 stt) [5] 11)1 is an -nomm, which is non- negative. Thus, d, (2,,2)) 2 0. Additionally, all three types of distances satisfy symmetry, hence da (2s, 2)) = dx (1),
Proof. Each term in the metric (z, 2,) is non-negative:
To derive the triangle inequality for latent task belief metric, we need to verify that for any three task beliefs 2,, 2j. and žg, the following holds:
dx (zi, 2x) S dx 23) + dx (2z, 2x)
Expand the distance (21, 2,) and du (zj, =k) (2,, 2,) R (s - RT (5} ag) + We
(st az) (5) 64) ,
and similarly, du (2j =x) = [R] (5] a ($ a*) ) Expand the distance (z,, =k)
og), (st an) a)))+1115 1 Ill
du (2,, =x) = R (s a,) R [ ax)) + (T, (s a.), (s a +1 (st s
1 (s $t
To prove the triangle inequality, we need to show that:
|RT a.)) - R2 (ot ax)l S |R (sT ..)) (xJ a;)) + |RJ (s)t an) - Rã (+ ax) This follows directly from the standard triangle inequality for absolute values. Similarly, for the Wasserstein distance term:
(2(d,) (s* a4),TT (st. ax)) S W2(dx) (5) (a,ls...)))WW/.. w)))) .ay), (T),TT((5 (x)))))) ())
This holds due to the triangle inequality for the Wasserstein distance (Villani et al. 2009). Finally, for the inverse dynamics term:
|14 ./+) [g[s (st stt Il1 S (s47(s (66)) [17((] (s) - [5 St)
This also follows directly from the standard triangle inequality for norms. Since each of the three components satisfies the triangle inequality, we can conclude that the latent task belief metried, also satisfies the triangle inequality:
dz (21, 2x) < dr (26, 23) + d,y (zz, 2x)
15

---

Published as a conference paper at ICLR 2025
B THEOREMS AND PROOFS
Lemma 1. Let S+ ((S) x Z be the latent dynamics space, and a is the improving policy. The F function is defined as follows:
(s) (2s, 23) = =I (sT az) - ag)) + (st,, (5) 0z))
(3) 18441) [5] 111,
(14)
where st (g(s:), 21), 5 (g(s)), 2n) the augmented states corresponding task and task j. and 24, 24 represent different task beliefs in latent space St. Then, there exists unique least fixed point such that F(d) d.
Proof. To prove that F has a unique least fixed point, we need to establish that F is monotonic and continuous, ensuring the conditions for least fixed-point result.
We first show that F is monotonic. Suppose S d', meaning d(zz, 23) S d' 23) for all 2i 2j in latent space. Then, by the definition of F, we have:
W2 (2(dx) (s,t a.), 17 aB)) S (2(d)) (s 13) az)).
Since the Wasserstein distance is non-decreasing in d, it follows that F(d, π) S F(d', T). Addition- ally, the reward difference and inverse dynamics term remain unchanged under ordering, confirming that F is indeed monotonic.
Next, we show that J continuous with respect to the pointwise supremum an w-chain (du)- Consider the sequence {dn} forming an increasing chain in M, meaning d, (2;, zj) converges to doo (Zz, 2,). By the properties of the Wasserstein metric and L1 norm, limits commute with these operations, leading to:
FL dn)( -186 (57 a.)) - R35 (1655aay))) W2( d.) (6), nEN
nEN
+ (177 -
|R (s , 0.) - (s(s :4) sup
sup (2dd .o). (")
nEN
+ (st stt) -
sup( (0) ..,) -
nEN
+ |1 (
(L) F (dw (2n,)))) neN
(15)
This guarantees that F is continuous. The policy in the above equations is fixed. As the policy updates, it can converge to a fixed point.
Since M complete partial order (w-CPO) and F is monotonic and continuous operator, the fixed-point theorem (Ferns al. 2004) ensures the existence of least fixed point d such that F(d) d.
Theorem (Value difference bound). Given two tasks Mi and Mj in the latent space with states S+, and let v be the value function of policy the value difference bound between the tasks can be given by:
E st.
(=(+)) d=(z,
(16)
where d (2i, 2z) is the latent task belief metric.
16

---

Published as a conference paper at ICLR 2025
Proof. We will use the standard value function update:
= +Y
P s.a))
P"
3'E8
with ViT = 0, and our task difference metric de (z,, and prove this by induction, showing that for all n E N and states ,,,88 E S+:
(T -17)) (+T)) (z1, The base case holds trivially:
-V = 0, so assume this holds for n.
Now, for + 1:
|V ( (85)) (
a.) s'Es+
ETST st at,a) (
-
Σ 17 (df)).
sEst
We decompose this into rewards and transition dynamics:
R (st (s) a)) a;))
Σ
(TT st - 1776(111)))))))0)))))
s'ESt
Using the Lipschitz property of the value function, we can bound the transition dynamics difference:
[ (st, a4) RJ (st a3)l + wW2 (7 (st, a.), (, (sJ ag)) + 1 -
By the definition of the latent task belief metric d, (2,, 2z). we know:
(z, 2j) RT (s a.) a:)- (5] aay) (T, a.), (7 az)) (T 1767 166
We assume that z contains sufficient information about the task differences at the current time step, and we disregard the effect of the discount factor y during the experiments. Thus, we have:
V
(s+) V (55) I d 2z)
Lemma 2. Let g be (CR, Cr, -apprrxximate bisimulation abstraction of M. The states in the real space are mapped to the latent dynamics space S+ through g. For any two states in the real space, 81 and 82 their representations in the latent space identical, i.e., g(s1) g(S2), then
[R(1,, R(s2, a)l S ER:
([T(s1,)) T(82, a)ll SET,
[((51,8)) [S2, 3 $ Ez-
Proof. The proof follows similar process to Lemma 3 in Tiang 2018). Here, have added representation constraint on the inverse dynamics.
Theorem 2 (Latent transfer bound). Let Q'm, be the optimal Q-function for task Mj. The difference between Qm, and the Q-function of the policy π learned from task M,, applied to task Mj. is bounded as follows:
Rmax
Qin, - -QThm SER ( ( 1| - -zyll)) 2(1 77)
(17)
17

---

Published as a conference paper at ICLR 2025
Proof. ||2 2y |la is the distance between the latent task beliefss of task M, and task Mj. In the real state space, the augmented state is still conditioned on 21 (Section 3.3 )
We aim to bound the difference between the Q-function QM,, (the optimal Q-function for task M2) and I,,, (the Q-function for policy л learned on task M, but applied task M,). This difference is given by:
IiMM, - IM,,,M,llo-
The Q-function for any task satisfies the Bellman equation:
0)) R(s, a) + (( max Q(s',a')
Using this, we decompose the Q-function difference:
1 1QM, 19M.J, - 1-9 1li, - -T,, loo
Here, Th, QM. RM, (st, a) + T.BywTTn (el (0) (dt,, ,d)].
According to Lemmal
sup a 13 Est g(x))=(())))
[TM a) (s).
(1Txi(st, a) (s a))1, + || (s).a) - (..))))1)))
sup of oj ESt (((=))=(=())))
sup 25 ESI, {(x))==(ss))
|su a) - (st
sup
TM, (8) - -T,, (5f aa)l:
st st.07 (
ET + [s] a) TM, a))
(18)
RR. (st,a) RM, (5)
sup of 10 Est g(e))==(()))
(st a)) (8 a)) + |RM, (8t - RM, (5), a)l
sup of oj Est, s(()))=()=))))
sup .5 ESI =((=))=((s))
[R, ()) - -Rm
sup
|RM, (8f .)) - RM, (s),a))
at 01.05E81 ((.))==((,))
= ER + RR, (5) RM, RM,
(19)
IMs -mn (st,,
sup st 1g Estt (
(mil () 10 ( a o3tt)1, + - (55t)))))))))
In,
sup
of oj est, {(=:)=g(es))
(M s85) Im
($) (5)t
sup
sup
.f sT..tSt. s(=))()))))=)())
ot at.5 Est (()))()))))())))
= =er + (M. -1 (m,
| 1,
(20)
[T a) - Tu, (5) a)]11 + (R RM, a)]+ Imilst, ((t) - 15, +b11)|
E + + || 2; |l,
(21)
18

---

Published as a conference paper at ICLR 2025
10M, - TM, llo
(TM (6) - -T((m))
[s] a) + (TM, (s)tt,)),,, "0) (i,) - (s - (
S $ent7
(TM, (8] ,0) vM,)) - (6 ,0), [Mn
Sct + (8) a), VM,) - -(Tmi .), (MM (6]1Mm(67.88) TM, (83.88887
Rmax 11100
Sen + 7(T +, + || zy]li)|
2(1 - 7) 2(1 77)
<en 7(T + +1 1 21|11) Rmax
We get the final bound:
Qin, - 19hiM
SER ter 2 -zylli) 2(1 7)
Rauxx
C SIMBELIEF PSEUDO-CODE
Algorithm SimBelief algorithm
Input: distribution p(M)
Initialise: context encoder do. real envs dynamics Pa, actor no. critic Qw. replay buffer D, belief similarity learner th, latent dynamics pe, state encoder g, distribution offset (Au, Ao)
1: while not done do
2: Sample tasks Mtrain (M.) from p(M)
3: Collect trajectories with Ty and add to buffer D
b Data collection Training step
4:
for step in SAC training steps do
5:
Sample training tasks Musin from p(M) Sample batches from D
6:
7:
Infer specific task beliefs 2, ~ thr (br h) h Ta) Infer latent task beliefs žι ~ (ht h) T.))
8:
9:
Permute zf to get zf
10:
5,- = (g(s), 21). s = (g(s).z,) 2,)
Update latent dynamics pe and state encoder g using Eq.
12:
Update belief similarity learner ei using Eq. Update distribution offset (Aμ,, Ao) using Eq.
14:
Integrate of with b; to obtain b*
15:
Update (0,w) with SAC algorithm end for
P SAC update
16:
17:
for step in VAE training steps do
18:
Sample T.r B with trajectory length T Decode only the past trajectories and update qas tb, and ps using Eq. 6 end for
19:
> VAE update
20:
21:
end while
D CONTEXT-BASED META-RL BASELINES
In this section, we provide a detailed overview of the baseline methods compared in our experiments, highlighting their methodologies and design principles.
PEARL (Rakelly et al. 2019) (Probabilistic Embeddings for actor-Critic RL) is an off-policy meta- reinforcement learning algorithm designed enhance both meta-training sample efficiency and
19

---

Published as a conference paper at ICLR 2025
rapid adaptation to new tasks. By disentangling task inference from control, PEARL employs probabilistic latent context variable that enables structured exploration and efficient posterior sam- pling. This design allows the policy reason about task uncertainty and adapt quickly in sparse reward or dynamic environments. Built upon the Soft Actor-Critic framework, PEARL. achieves better sample efficiency compared on-policy meta-RI. methods.
MetaCURE (Zhang et al. 2021c) (Meta-RL with Efficient Uncertainty Reduction Exploration) an off- off-policy meta-RL framework designed to address the challenges of sparse-reward environments. explicitly separates exploration and exploitation by learning distinct policies for each, enhancing the efficiency of task inference and adaptation. The exploration policy driven by an empowerment- based intrinsic reward that maximizes information gain about the task, enabling efficient collection of task-relevant experiences. MetaCURE. leverages a shared probabilistic task inference mechanism, which improves sample efficiency by integrating exploration and exploitation processes. Compared MetaCURE. SimBelief does not require knowledge of task IDs during the training phase. Instead, it learns the common structure of tasks to enable reasoning, resulting in stronger online adaptation capabilities.
VariBAD (Zintgraf al. 2019) (Variational Bayes-Adaptive Deep RL) is a meta-reinforcemeen learning framework that approximates Bayes-optimal policies by leveraging variational inference and latent task embeddings. The algorithm learns posterior belief over tasks using a variational auto-encoder and conditions its policy on this belief to balance exploration and exploitation in un- certain environments. VariBAD is notable for its ability to perform structured online exploration by integrating task uncertainty directly into action selection. Unlike traditional methods that rely on computationally intractable posterior sampling or explicit planning, VariBAD offers tractable and flexible approach to Bayes-adaptive policies.
HyperX (Zintgraf et al.) 2021) (Hyper-State Exploration) is meta-reinforcement learning method that addresses sparse reward environments by leveraging exploration bonuses to meta-learn approxi- mately Bayes-optimal aask-adaptation strategies. It integrates two exploration bonuses during meta- training: (1) a bonus based on hyper-states (combining environment states and task beliefs) to en- courage diverse task-exploration strategies, and (2) a reconstruction error bonus to incentivize the agent collect data where task belief's are inaccurate. By exploring hyper-states, HyperX efficiently gathers data for belief inference and optimally trades off exploration and exploitation in sparse or complex environments.
RI. (Duan et al. 2016) reformulates the reinforcement learning process as a meta-learning problem, embedding task-specific learning within the hidden state of a recurrent neural network (RNN). By leveraging "slow" reinforcement learning algorithm to optimize the RNN's weights, RL enables the network to store and process information about task dynamics across episodes. This design al- lows the agent to efficiently adapt unseen tasks by utilizing its historical trajectory data, including observations, actions, rewards, and termination flags. RL demonstrates competitive performance in solving multi-armed bandits and tabular MDPs, achieving results comparable to theoretically op- timal algorithms. Additionally, its scalability to high-dimensional tasks, such as visual navigation in dynamic environments, highlights its potential as versatile and efficient meta-RL framework.
E ENVIRONMENTS AND IMPLEMENTATION DETAILS
Table 1: Adaptation length and goal settings for environments used for evaluation
# of adaptation
Max steps per episode
Environment
Goal type
Goal range
Goal radius
episodes
Cheetah-Vel-Spare
200
Velocity
[0.3] Semicircle with radius
0.5
Point-Robot-Spee Robot-Sparse
60
Position
0.3
Walker-Rand-Paee Rand--aanndnn
200
Velocity
1.5
0.5
Panda-Reaad Reach
50
Position
0.05
Panda-Push
50
Position
0.05
Pick-And- And-Place
50
Position
0.05
20

---

Published as a conference paper at ICLR 2025
Panda-Reach Reach
Panda-Puse Push
Panda-Pice Pick-And-Plan
Panda-gym tasks. The Panda-Reade task involves controlling the Panda robotic arm to move its gripper to randomly generated target position within defined workspace. the Panda- Push task requires pushing a cube from its initial position to randomly generated target position on table, and the Panda-Pick Pick-And-Plaatt task focuses picking up cube and placing at randomly generated target location above the table.
Figure 7:
Table summarizes key settings for the environments used in this work, including the number of adaptation episodes, the maximum number of steps allowed per episode, the type of goal (whether velocity or position-based), the goal range, and the goal radius. The specific environments are described as follows:
Cheetah- Cheetah-Vel-Sparse Vel-Sparse. The Cheetah- Vel-Sarse environment involves controlling a half-cheetah robot to achieve and maintain target velocity, which is randomly sampled from a uniform distribu- tion over the range [0. 31. At each time step, the robot receives a reward based on how closely its current velocity matches the target velocity, with sparse reward system in place. Specifically, if the difference between the current velocity and the target velocity is within a specified tolerance (goal radius of 0.5), the robot receives reward; otherwise, the reward zero. Additionally, control cost applied to penalize large actions. The observation space includes the robot's joint positions, velocities, and body orientation, while the action space consists of motor torques controlling the robot's movements. The sparse reward encourages the agent to achieve the target velocity while minimizing control effort.
The Point-Robot- Saboot-ppase environment requires a point robot to navigate to- wards goal randomly placed along a unit half-cirele. The goal position sampled at the beginning of each episode, and the robot receives rewards only when it reaches within goal radius of 0.3 units from the target. The observation space consists of the robot's current position, and the action space allows movement in both and directions. The reward system is sparse, meaning that rewards are only given when the robot within the specified goal radius, encouraging the robot to explore and move efficiently towards the target. The sparse reward structure makes the task more challenging, as the robot receives feedback only when it approaches the target, requiring it to learn how to navigate effectively with limited guidance.
Point-Rove SonttRRooot-Sparee
The Seiteer-Rand-Parte environment involves controlling a bipedal walker robot, with the added complexity of randomized physical parameters such body mass, leg strength, and joint properties. These parameters are randomized at the start of each episode, requir- ing the robot to adapt various physical configurations order move forward. The observation space includes the walker's joint angles, positions, and velocities, while the action space consists of motor commands applied the walker's joints. The reward system is primarily based on how closely the walker's 's velocity matches a target velocity of 1.5 units per time step. If the deviation from the target velocity exceeds 0.5 units, the reward is set O. For smaller deviations, the reward is given as 0.8 distance from the target velocity. Additionally, small control cost proportional to the sum of the squared motor actions is subtracted from the reward penalize large actions.
Walker-Ramn
Pench-Reach. The Panda-Reach task involves controlling the Panda robotic arm's gripper to reach a target position, randomly generated within a 30cm x 30cm x 30cm workspace. The task considered successful when the distance between the gripper and the target is less than 5 cm. The observation space for this task includes the position and velocity of the gripper (6 coordinates), augmented by two additional vectors representing the desired goal (target position) and the achieved goal (current gripper position). The action space comprises three movement coordinates (x, z) for controlling the gripper. The gripper remains closed throughout the task, and its movement along
21

---

Published as a conference paper at ICLR 2025
these axes determines the success of the task. A sparse reward function is used, with reward of 0 the target is reached, and -1 otherwise.
Panda-Push. The Panda-Pusk Push task requires the robot to push cube (side length of cm) placed on a table to a target position. Both the target position and the initial position of the cube are randomly generated within 20cm 20cm around the neutral position of the robot. The gripper remains closed during the task, and the objective to move the cube to within 5 cm of the target position. The observation space includes the gripper's position and velocity (6 coordinates), along with the cube's position, orientation, and velocity (12 coordinates). The action space consists of three movement coordinates (x, y, z) for controlling the gripper's movement. As in the Panda- Mead- task, a sparse reward function is used, providing a reward of O if the task is completed successfully, and 1 otherwise.
task involves the robot picking up cube (side length of 4 cm) and placing it at target location, which is randomly generated within 20cm x 20cm 10cm volume above the table. The task is completed when the cube is placed within 5 cm of the target position. The observation space for this task includes the gripper's position and velocity (6 coordinates), the cube's position, orientation, and velocity (12 coordinates), and the opening state of the gripper (1 coordinate). The action space is expanded to include three movement coordinates (x, y, z) for controlling the gripper and an additional coordinate for controlling the gripper's 8 opening and closing. A sparse reward function is employed, rewarding the robot with 0 for completing the task and 1 otherwise (Figure 7
Panda-Pick-AAddPPlaee
The Panda-Pick-Ande
All experiments were conducted using an Nvidia RTX 4090 GPU, the source code is available at: https://githin com/
Table 2: Hyperparameter settings for SimBelief in different environments
Parameter
Cheetah- Cheetah-Vel-Sparse
Panda-Push
Sontt-Rooot-Sparee
Panda-Reach
Name
Walker-Rand-Pe
Number of Tasks
120
100
100
60
Number Training Tasks
100
80
80
50
Number Evaluation Tasks
20
20
20
10
Number of Episodes
2
2
3
3
Number of Iterations
1000
2000
1000
4000
RL Updates per Iteration
2000
1000
1000
1000
Batch Size
256
256
256
256
Policy Buffer Size
le6
le6
le6
le6
VAE Buffer Size
lc5
5e4
5c4
5c4
Policy Layers
[128, 128, 128]
[128, 128]
[128, 128)
[128, 128, 128)
Actor Learning Rate
0.0003
0.00007
0.00007
0.00007
Critic Learning Rate
0.0003
0.00007
0.00007
0.00007
Discount Factorr (?)
0.99
0.9
0.9
0.9
Entropy Alpha
0,2
0.01
0.01
0.01
VAE Updates per Iteration
20
25
25
VAE Leurning Rate
0.0003
0.001
0.001
0.001
KL Weight
1.0
0.1
0.1
0.1
Task Embedding Size
10
10
5
5
F THE MECHANISM OF SIMBELIEF FOR RAPID ADAPTATION
In this section, we will explain, from the perspective of task belief representation, why SimBelief can quickly converge to near Bayes-optimal policy during the training phase and how it enables adaptation to OOD tasks within a single episode sparse reward environments.
F.1 THE EVOLUTION OF TASK BELIEF DURING THE TRAINING PHASE
During the training phase, the agent needs to explore complex environments to acquire relevant information about the current task, which is represented by the specific task belief b,. Other meta- RL algorithms are limited to task-specific information while ignoring the similarities between tasks. SimBelief's ' latent dynamics can capture the unique structure shared among tasks and distinguish these structures through the latent task belief by Specifically, during the learning process, SimBelief
22

---

Published as a conference paper at ICLR 2025
Early Stage Training
Middle Stage of Training
The Stage Training Convergence
Cheeton vel
Dersity
Value
Value
Early Stage of Training
Middle Stage Training
The Stage of Training Convergence
Deriiiy
Value
Value
Value
Early Stage of Training
Middle Stage of Training
The Stage of training Convergence
Desicity
Velue
Value
Value
Early Stage of Training
Mi dilee Stage of Training
The Stage of Training Convergence
Finda
Valun
Wilue
Specific task belef
Lacont task belief
Mixture distrbbttin
Figure 8: Visualization of different task belief's during the training phase.
identifies the distribution of similar tasks via by and then uses b, to the agent to the exact specific task distribution, illustrated in Figure [ In Figure g we visualize the specific task belief, latent task belief, and their mixed Gaussian distribution. During the early and middle stages of training, the agent quickly identifies the distribution of the task being executed. After convergence, it can accurately fit the specific task distribution.
F.2 THE CORRELATION BELIEFS BETWEEN DIFFERENT TASKS DURING THE ADAPTATION PHASE
We visualize the correlations of task beliefs across different tasks in Reach, Panda-Push, and Panda-Pick-Place environments. each environment, 20 tasks are ran- domly generated, and the SimBelief agent's rollouts are used to extract the specific task beliefs and latent task beliefs. The cosine similarity between task beliefs across tasks is calculated, and the correlation matrices visualized. As shown in Figure 9 specific task belief primarily captures the local information task distributions (learning weaker correlations between tasks), while latent task belief emphasizes global information (capturing stronger inter-task correlations). This illus- trates the principle behind SimBelief's s ability to achieve rapid adaptation to OOD tasks in sparse reward environments: latent task belief enhances the agent's reasoning efficiency across tasks and
Panda-
23

---

Published as a conference paper at ICLR 2025
improves the efficiency of knowledge transfer. Theorem demonstrates the transferability of tasks in the latent space.
G ABLATION STUDY
G.1 LATENT DYNAMICS ABLATION
To validate the effectiveness of algorithm using the learned latent task belief, froze the overall latent dynamics and, similar to VariBAD (Zintgraf et al. 2019), only used an external VAE for environment reconstruction. Our algorithm demonstrated significant improvements across all six tested tasks, with particularly notable gains in the more exploration-demaning tasks, Panda-Push and Panda-Pick- Piiila-Pic-AAn---- Fick-And-PPlae (Figure 10
G.2 THE IMPACT OF INVERSE DYNAMICS ON ADAPTATION ABILITY
To demonstrate the effectiveness of the inverse dynamics module in facilitating rapid learning of task similarity, we conducted ablation experiments by removing the inverse dynamics module from the latent space Cheetah-Vel-Spane Vel-Sparse (velocity range [4.0,5.0] and Point-Robot-Spenn tasks. The results show that the inverse dynamics module plays a critical role in enabling the agent's reasoning and generalization on OOD tasks (Figure IT..
(radius=1.2)
G.3 ABLATION STUDY W, AND Wi
When combining by and by. we apply mixture of Gaussians to the two beliefs. To evaluate the impact different weights, we compared three weighting configurations on the Cheetah- Sparse task: [w,, wy] [0.5, 0.5]. [w,, wy] [0.25, 0.75], and [w,, [0.75, 0.25]. Among these, [w,,, wi] = 00.5, 0.5) achieved higher exploration efficiency during training and demonstrated stronger adaptability to OOD tasks (Figure 12)
G.4 ABLATION STUDY OFFSET (Aμ, Ao)
During the training phase, we use the offset (Aμ, Ao) applied to the latent task belief primarily improve the stability of the algorithm's 's convergence in continuous control tasks (e.g.. Cheetah- Vel-Sparse). However, this does not have significant impact the overall performance of the algorithm. Even without using the offset during training, SimBelief still outperforms other baselines (Figure 13
H ADDITIONAL VISUALIZATIONS
In this section, we provide the task belief representations of SimBelief and other baselines on the Galker-RanndPaarnn tasks (Figure I4), the adaptation performance on in-distribution tasks Panda- Pick-And-Place and Walker-Rand Params (Figure 15). and the exploration performance on the in- distribution Point-Robe SonttRobot-Soe Sparse task (Figure 16
24

---

Published as a conference paper at ICLR 2025
Specific Task Belief Correlation
Latent Task Belief Correlation
Point-Robot-S Spaat- Sparse
Specific Task Belief Correlation
Latent Task Belief Correlation
Panda- Panda-Reach Reach
Specific Task Belief Correlation
Latent Task Belief Correlation
Panda- Bunda-Push Push
Specific Task Belief Correlation
Latent Task Belief Correlation
Panda- Pick. Place
Figure 9: The correlation matrix of task beliefs across different tasks. Specific task belief focuses more on local information between tasks, while latent task belief captures the global characteristics of task distributions.
25

---

Published as a conference paper at ICLR 2025
Cheelsh--Ve-Sparee
Point- Pant-Rabbol-Speeer
Valla--anns-aaaam
76
Aseraga Ratum
800
120
2
0.5
1G
15
20
Frames
Framcs
Paidia Reach
Panda-Push
Parda Piike And--laoe
1.0
St00aaR Rugl
SUNOWN Pote:
0.0
0.4
84
0.2
0.
00
1.0
Framsri
Trames
SimBelief
SimBelicf w/o Latent Task Bellef
Figure 10: Ablation study on SimBlief's latent task belief.
Sharaah-Vel-Sparse
Point-Robot-Sparse
Average Retur
Average
15
-10
10
-15
Episodes
Episodes
SimBeliet
SimBeliet w/o inverse dynamic
Figure 11: Ablation study on SimBlief's inverse dynamic module in latent space.
Choorsh Vell Sparse
Velocit Range (0.C 3.
Velccty Range: p... 4.00
Valoaly Fange: H.O. 5.2
-0 300
50
Troves
Splcodes
Calecdes
Splecdes
Figure 12: The performance of different Gaussian mixture weights on the Cheetah- Sel-Spasee task.
Checlat-
Voliccty
155
Fiares
sode
Epiche
SinBoled
Beldellel sya offere
Figure 13: Ablation study on offset (Au, Ao).
26

---

Published as a conference paper at ICLR 2025
let spedifsc
Sindellr latent
Muacunn
MISS
Figure 14:
The task belief representation of all experimental algorithms for 10 random in- distribution Walker-Ranc tasks.
Panda Pick And Place
Walker Rand Parama
-25
Average Return
Awerage Return
70
-35
60
-40
-50
Episodes
Episodes
SimBeleff
PEARL
MetaCURE
VariBAD
Hyperk
Figure 15: Average test performance for the first 5 rollouts on randomly generated tasks of Panda- Fick-And-Place and Saikker-Rand-Parce
SimBelief
MetaCURE
Episcode
Episode
Figure 16: Exploration and adaptation performance in the random in-distribution (radius = 1.0) tasks of Roint-Robot-Sparse.
27


Reviews:
[{"review_id": "FAotWqRfrL", "reviewer": "Reviewer_6WZh", "summary": "The paper uses the bisimulation metrics to measure task similarity and learn the common structure across tasks for rapid task identification and adaptation in meta-RL.", "questions": "- No major flaws with the paper. The key difference and advantages over VariBAD could be further enhanced and elaborated, since both infer task beliefs using the world model.\n- Some notations are not well explained and kind of confusing, e.g., $z_l$ and $z_r$ are confusing, and why use two distributions?\n- The proposed method is motivated by the bisimulation metric, which is originally proposed for deterministic MDPs (Castro, 2020). Can the proposed method scale to stochastic MDPs?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The proposed method is well formulated and clearly presented.\n- The effectiveness of the latent task belief metric is validated by theoretical guarantee. \n- Experiments demonstrate the superiority of the proposed method over strong baselines, especially the generalization capabilities to OOD testing tasks.", "weaknesses": "- The topic of online meta-RL is kind of old. \n- The core of the proposed method is using the reward and state transition functions $p(s’,r|s,a)$, or called world model, to measure task similarity in a latent space and hence infer task belief for context-based meta-RL. This kind of task inference has been investigated by existing works like VariBAD. \n- The baselines are kind of old, mainly 2019-2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PVHWyYfg1A", "reviewer": "Reviewer_RXre", "summary": "This paper presents a novel framework called SimBelief, which effectively extracts common features of similar tasks using the Bisimulation metric for meta-RL tasks.", "questions": "(1) Regarding readability: In Sec. 2.1., N is not properly defined before use. Inverse dynamics in Definition 2 or Eq.(3) is used before it is properly defined.\n\n(2) Some of the manuscript's functions are unclear. What is the output of inverse dynamics $I_{i}^{\\pi}(s_i^{+},s_i^{',+})$ ? \nWhat is the input and output for $q_{\\phi}$, $\\psi_{l}$, $\\psi_{r}$ ?\n\n(3) Similarly, In Algorithm 1, it is unclear $b_{r}^i={q_{\\phi}(z_r^i|\\tau_{:t})}_{0:T-1}$ means. \n\n(4) What are the $L_{rec}$ and $L_{bisim}$ in Figure 2? If they are defined in the manuscript differently, please use the same notation.\n\n(5) How to determine a proper $w_r$? and how sensitive is it to the overall framework?\n\n(6) The proposed method is complex and contains various components, so presenting only a simple algorithm raises questions about reproducibility. Do the authors plan to release the code for the benefit of the RL community?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper is generally well-constructed and well-motivated and provides some theoretical background on the proposed method.\nIn addition, the paper has the potential to help the community understand the use of latent embedding for task generalization.", "weaknesses": "The paper includes many symbolic notations, which can confuse readers without strict and coherent representation. For example, in Fig 2, it seems $q_{\\phi}$ outputs h, but the notation of Eq. (6) or others in the manuscript, $q_{\\phi}$ outputs $z_r$.\n\nSimilarly, some of the notations are used before proper definition, which hinders the readers from fully understanding the contents.\n\nThe paper's reproducibility is questionable as it consists of various complex components for the framework.", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XCereFztQt", "reviewer": "Reviewer_cdUU", "summary": "The authors propose a novel meta-RL method tackling the problem of efficient task identification in sparse reward settings, where methods that do reward reconstruction are not sufficient.  In a context-based meta-RL framework, they propose inferring the latent task representation through a Gaussian mixture of a variational latent representation and a new task belief similarity latent representation the authors introduce.  This task belief similarity latent is trained to model a  bisimulation-inspired latent task belief metric between two tasks.  They demonstrate modest improvements compared to prior methods in sparse reward, simulated locomotion and manipulation tasks and robustness to OOD task variations.", "questions": "* Why is the inverse dynamics required in Definition 2?  Have you done any ablations for how it compares to the bisimulation metric?\n* What are the offsets mentioned in section 3.3 and how are they trained?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "* Introducing ideas of modeling task similarity to meta-RL is interesting and worth exploring\n* Sound experimental setup and results", "weaknesses": "The biggest weakness is that the proposed method lacks motivation and reasoning about how they achieve the claims the authors make.  \n* Combining the variational task belief $z_r$ with their proposed task belief similarity $z_l$ through a Gaussian mixture.  This does not combine both types of task representations, but instead samples from one or the other.  Furthermore, $z_r$ already models dynamics information in order to reconstruct the trajectory, so it's unclear what information $z_l$ adds. This is also flawed because the task belief similarity is not modeled as a distribution (at least based on Section 3.2).\n* There are a couple things in Section 3.3 that are mentioned briefly but not explained: using the Q-function to train offsets for the task similarity distributions and minimizing KL divergence to the variational task belief $z_r$ when predicting $z_l$.  These details seem important to the method and the fact that the KL divergence is required to “ensure the agent does not confuse similar tasks” indicate that $z_l$ may not be as effective as the authors claim.\n\nOther major weaknesses include:\n* Clarity of writing.  The method section, especially 3.2, was difficult to understand.  The experiments section does not describe what the different tasks are in each environment, and Figures 4 and 5 are hard to interpret.\n* Insufficient discussion of relevant related work, especially MetaCURE, which seems to tackle the exact problem of task ID with sparse rewards.\n* Insufficient discussion of results, especially providing insight into why comparison methods may under or over perform.\n* No ablations provided.", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "B9nzRAPChr", "reviewer": "Reviewer_z5au", "summary": "This work considers meta RL, especially these sparse reward settings. This work proposes to utilize bisimulation metrics to extract behavioral similarity in continuous MDPs and proposes SimBelief, which can extract common features of similar task distributions. This work further theoretically validates the effectiveness of the latent task belief metric in BAMDPs. Experiments show the effectiveness of SimBelief, especially in sparse reward settings and o.o.d settings.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- It is novel to include Bisimulation to measure the task similarity.\n\n- Sparse reward is a hot topic in the meta RL community.\n\n- Several experiments show that SimBelief performs better than several baselines, especially in sparse reward settings.", "weaknesses": "- The writings of the proofs are poor with several typos. \n\n(1) in lines 714-715, the value function update should be $V_{n+1} = ... V_n$\n\n(2) in lines 719-710, $d^n$ here seems to be $d$\n\n(3) in lines 733-739, why we can use the Lipschitz property of the value function to get the result? It seems the core to prove the result but the proof does not include a detailed discussion of this inequality.\n\n(4) If you use the Lipschitz property of the value function, what is the Lipschitz constant?\n\n(5) The proof of Theorem 2 is even more poor and needs to be thoroughly polished.\n\n- The baselines compared in the paper are somehow old. The latest baseline in this paper is MetaCURE and  HyperX, which were published in 2021. There are several much more new baselines for meta RL, especially sparse reward settings or o.o.d. settings like [1-3], which should be discussed and compared. Also, it seems the evaluation environments are designed by this paper, are these previous works utilizing these environments?\n\n- Sparse reward is indeed an important setting in RL, but why the proposed method can boost the agent's performance in sparse reward environments? As the reward structure may be independent of the dynamic structure, \"extracting common features of similar task distributions\" may not benefit the extremely sparse reward settings. Assume that there are n tasks that own the same dynamic, and we can only get the reward in the terminal step, there is no way to identify what is the current task as it is independent of the dynamic structure / previous rewards. Although it is not my major concern (experiments show that the proposed method can perform better than baselines in sparse settings), it will make the paper more solid if there are theoretical analyses about why the proposed method can handle sparse reward settings.\n\nRef:\n\n[1] Learning Action Translator for Meta Reinforcement Learning on Sparse-Reward Tasks\n\n[2] Meta-Reinforcement Learning Based on Self-Supervised Task Representation Learning\n\n[3] Enhanced Meta Reinforcement Learning using Demonstrations in Sparse Reward Environments\n\n\n------\n\n**After rebuttal, most of my concerns have been addressed, especially about the theoretical proofs, and I have raised my score to 6.**", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)