PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
Authors: Dongchi Huang
Jiaqi WANG
Yang Li
Chunhe Xia
Tianle Zhang
Kaige Zhang
Venue: ICML 2025 poster (2025)
Keywords: Reinforcement Learning; World Models; Safe Reinforcement Learning; Model-based Reinforcement Learning; Privileged Learning
Primary area: reinforcement_learning->deep_rl
TL;DR: An algorithm that leverages privileged information to address partial observability in safe reinforcement learning.

Abstract: Partial observability presents a significant challenge for Safe Reinforcement Learning (Safe RL), as it impedes the identification of potential risks and rewards. Leveraging specific types of privileged information during training to mitigate the effects of partial observability has yielded notable empirical successes. In this paper, we propose Asymmetric Constrained Partially Observable Markov Decision Processes (ACPOMDPs) to theoretically examine the advantages of incorporating privileged information in Safe RL. Building upon ACPOMDPs, we propose the Privileged Information Guided Dreamer (\textit{PIGDreamer}), a model-based RL approach that leverages privileged information to enhance the agent's safety and performance through privileged representation alignment and an asymmetric actor-critic structure.
Our empirical results demonstrate that \textit{PIGDreamer} significantly outperforms existing Safe RL methods. Furthermore, compared to alternative privileged RL methods, our approach exhibits enhanced performance, robustness, and efficiency.
Codes are available at: https://github.com/hggforget/PIGDreamer.

PIGDreamer: Privileged Information Guided World Models for Safe Partially
Observable Reinforcement Learning
Dongchi Huang'
Jiaqi Wang² Yang Li
Chunhe Xia
Tianle Zhang? *Kaige Zhang*
Abstract
Markov decision processes (CMDPs) (Altman, 2021), in which agents operate on the underlying states with the objec- tive of maximizing rewards while ensuring that costs remain below predefined constraint thresholds. However, existing approaches overlook the inherent partial observability in real-world applications, thereby leading to the suboptimal deployment of Safe RI algorithms.
Partial observability presents a significant chal- lenge for Safe Reinforcement Learning (Safe RL), as it impedes the identification of potential risks and rewards. Leveraging specific types of privi- leged information during training to mitigate the effects of partial observability has yielded no- table empirical successes. In this paper, we pro- pose Asymmetric Constrained Partially Observ- able Markov Decision Processes (ACPOMDPs) to theoretically examine the advantages of incor- porating privileged information in Safe RL. Build- ing upon ACPOMDPs, we propose the Privileged Information Guided Dreamer (PIGDreamer), a model-based RL approach that leverages privi- leged information to enhance the agent's safety and performance through privileged representa- tion alignment and an asymmetric actor-critic structure. Our empirical results demonstrate that PIGDreamer significantly outperforms existing Safe RL methods. Furthermore, compared al- ternative privileged RL methods, approach exhibits enhanced performance, robustness, and efficiency. Codes are available at: https:// g thub. com/ Phogforget/PIGGDeaama
To mitigate partial observability in Safe RL, recent stud- ies (Hogewind et al., 2022; As et al., 2022; Huang et al., 2024) have investigated the application of world models (LeCun & Courant, 2022), resulting in remarkable success. These methods are situated within a more realistic frame- work: Constrained Partially Observable Markov Decision Processes (CPOMDPs) (Lee et al., 2018), where the agent operates on history of previous observations and actions, without direct access to the underlying states. Thus, world models (Hafner et al., 2023) which learn environmental dy- namics and task-specific predictions from past observations and actions, have been shown to effectively memorize histor- ical information and enhance sample efficiency. However, despite enabling agents to utilize historical information, these models introduce computational (Papadimitriou & Tsitsiklis, 1987) and statistical (Krishnamurthy et al., 2016; Jin et al., 2020) chanllengs under partial observations.
To address these challenges, recent studies (Hu et al., 2024; Lambrechts et al., 2024) investigate methods for leveraging privileged information within world models. The utilization of privileged information is a practical solution, in the real-world deployment of RL, only subset of sensors is accessible during testing (Hu et al., 2024) while a greater number of sensors is available during training. Furthermore, for agents trained in simulators and subsequently transferred to the real world through Sim2Real approaches (Yamada et al., 2023; Pinto al., 2017). the underlying states in the simulators be effectively leveraged to enhance the agents. Despite their remarkable empirical successes, these approaches still encounter two challenges: (1) The theo- retical understanding of leveraging privileged information in the world models remain relatively limited (Cai et al., 2025); (2) These approaches exhibit limited efficiency in leveraging privileged information (Hu al., 2024). These issues raise a crucial question:
1. Introduction
Safety presents a significant challenge to the real-world applications of reinforcement learning (RL) (Feng et al., 2023; Ji et al., 2024; OpenAl et al., 2019; Papadimitriou Tsitsiklis, 1987). Several researches (Liu et al., 2021; Thomas et al., 2022; Achiam et al., 2017) has been dedi- cated addressing this challenge, relying on constrained
"Equal contribution 'School of Computer Science, Univer- sity of Beihhang Beijing, China "School of Computer Science, Chinese University of Hong Kong. Hongkong, China 'JD Ex- plore Academy, Beijing, China "North Automatic Control Insti- tute, Taiyuan, China. Correspondence to: Kaige Zhang <zk- gusu@ gusu@gmail.com>.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforeement Learning
How can we achieve safe, high-performance, theoretically guaranteed and easily trainable agents under conditions of partial observability by leveraging privileged information?
As result, the agent must maintain belief state by, where dy (5) = Pr(s; s|he, bo) represents the proba- bility distribution over possible states given the history h fzu, ag. 21, A...... at of past actions and ob- servations, and the initial belief state bo- We denote the set consisting of all possible belief states as B, the belief reward function as RB (b, a) = Exes R(s) a), the transition function as 7(b, a, 2). For simplicity, we write T(b, a, 2) as baiz_ Finally, let the policy be denoted wo (a b). The objective of the policy is to maximize the long-term belief expected reward VR (bo) as follows:
To address these problems, we propose a novel frame- work Asymmetric Constrained Partially Observable Markov Decision Processes (ACPOMDPs), elucidating the mecha- nisms for integrating privileged information. Our theoreti- cal analysis within this framework demonstrates that agents will underestimate potential risks in the absence of privi- leged information. Furthermore, building upon our frame- work, we present Privileged Information Guided Dreamer (P IIGDreamer)) r), which enhances the world models, pre- dictors, representations, and crities with privileged infor- mation generally (i.e., the underlying states, past actions, and proprioceptive sensors). We integrate PIGDreamer with the Lagrangian method (Nocedal & Wright, 2006; Li et al., 2021) to ensure the Experiments on the satisfaction of safety constraints. (Ji et al., 2023b) and Guard (Zhao et al., 2024) benchmarks demonstrated that PIGDreamez surpasses existing Safe RL methods in both safety and performance. Furthermore, TGGDreames achieves 136% improvement in performance with only 28% increase in training time, illustrating superior uti- lization of privileged information compared to alternative privileged RL methods.
max VR (bo) = Eai= [Ei Ry (be, ax) [00], (I)
Constrained POMDPs (CPOMDPs) is a generalization of POMDPs defined by (S, P, R, Z, O, d. 7). The cost function set C = {(G,, b,)) comprises individual cost functions C, and their corresponding cost thresholds b,. The goal is to compute an optimal policy that maximizes the long-term belief expected reward VR (bo) while bounding the long-term belief expected costs Vc, (bo) as follows:
max VR (bb) = Eaiw 0' RB at) [ul,
V.t. (bo) = Ear 1 C'nnn (be, at) [b0] s b., Vi E [m]]
(2)
where in practical implementations, VR (bo) is updated by the following Bellman optimal equation:
Our key contributions are summarized as follows:
We introduce a novel theoretical framework, ACPOMDPs, wherein we theoretically demonstrate that asymmetric in- puts decrease the number of critic updates and facilitate the development of a more optimal policy.
Vi(() = -mxx aEA RB a) + Σ Pr(elb,a)V7(600))
=E7
and the Ve, (bo) is updated equivalently.
Built upon ACPOMDPs, we present P e r, a safe, negorperformance, and efficient algorithm that leverages privileged information in model-based RL.
Safe RL with Model-based Methods Model-based RL approaches (Moerland et al., 2022; Polydoros & Nalpan- tidis, 2017: Berkenkamp et al., 2017) provide significant advantages for addressing Safe RL problems by facilitat- ing the modeling of environmental dynamics. Jayant & Bhatnagar (2022); Thomas et al. (2022) enhance the inte- gration of model-free Safe RL algorithms with dynamic models by employing ensemble Gaussian models. Altema- tively, Zanger al. (2021) utilize neural networks (NNs) to quantify model uncertainty and constrain this error mea- sure through constrained model-based policy optimization. Recently, As et al. (2022) has integrated Bayesian meth- ods with the Dreamer framework (Hafner et al., 2020) to quantify uncertainty in the estimated model, employing the Lagrangian method to incorporate safety constraints. How- ever, these works overlook the issue of partial observability, which presents a more realistic setting in real-worle applica- tions. In this context, Hogewind et al. (2022) addresses the problem of partial observability in Safe RL by integrating the Lagrangian mechanism into the SLAC framework (Lee et al., 2020). Huang et al. (2024) achieved zero-cost per- formance by integrating the Lagrangian method with the
Our empirical results across diverse benchmarks indicate that PI PIGDreamer surpasses both existing safe RL meth- ods and privileged model-based RL methods.
2. Preliminaries
Sequential decision making problems under partial observations are typically formulated as Partially Ob- servable Markov Decision Processes (POMDPs) (Egorov et al., 2017), represented as the tuple (S, A, P, R, Z, 0) The state space denoted as S and the action space as A. The transition probability function is denoted as P Ps'' (s'|s, s, a). Let Z be the observation space, O (z|s', a) stands for the observation probability. The reward function is represented by R (s, a). The discount factor is represented by 7. In POMDPs framework, the agent has access only to the obser- vations 2n and actions a, at each time step t, without direct knowledge of the underlying state of the environment.
CPOMDPs
2

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforeement Learning
Dreamerv3 framework. From the perspective of POMDPs (Kaelbling al., 1998), constructing world models solely from partial observations does not fully exploit the potential of these models. In this work, we aim to address partial observability by leveraging privileged information that is ac- cessible during training, thereby ensuring the agent's 's safety and enhancing its performance.
long-term belief expected reward VR (0) while bounding the long-term belief expected costs Vo, (b):
n arg max (R(6). atA
(4)
s.t.Vc, (b) S b Vi E [m]]
Benefiting from the availability of the underlying states, at each time step, the critic receives and the action a and the underlying state 8, and updates the VR(8) using the following equation:
Privileged RL with World Models While extensively studied in the literature on model-free RL (Pinto et al., 2017: Salter et al., 2021; Baisero Amato, 2021; Baisero al., 2022; Li al., 2020), leveraging privileged information is seldom addressed in model-based RL. Yamada et al. (2023) represent the first attempt to utilize privileged information in the training of world models, employing it by distilling the learned latent dynamics model from the teacher to the student world model. Since the underlying state transitions of privileged information and partial observations differ, di- rect model distillation may eliminate essential features from partial observations. Lambrechts al. (2024) incorporate privileged information solely by introducing an auxiliary ob- jective that predicts the privileged information exclusively. Inspired by scaffolding teaching mechanisms in psychology. Hu et al. (2024) developed an approach to exploit privileged sensing in critics, world models, reward estimators, and other auxiliary components used exclusively during training, achieving outstanding empirical results. However, this ap- proach adds too many components, resulting in low training efficiency. Avalos et al. (2024) presents a method for approx- imating belief updates, supported by theoretical assurances that the resulting beliefs enhance the learning process of the optimal value function. In summary, the current exploration of world models that utilize privileged information remains aim to identify the optimal model architecture that efficiently and robustly incorporates privileged information within world models.
(i(() aeA R(s,a)+ + Σ P(oT)s (5)
max
Consequently, the VR(()) and Vo, (b) in the optimization problem (4) are estimated using this updated Vii(s):
Vi(0) 1000Emonp VZ)))
(6)
#ES
Notice that the update of the Vc, (b) is not presented, which is equivalent to (6).
3.2. Comparison with CPOMDPs
We compare the different estimations of the Va(b) and (b) to demonstrate the superiority of ACPOMDPs. Since the (b)(b) and (b) are equivalent in their estimations, we Collectively refer to them as V(b). For clarity, we rewrite the (b) in CPOMDPs as Vbym(b) and (bb) in ACPOMDPs (bs
Lemma 3.1. Kaelbling et al. (1996) showed that the value function time step I can be expressed by a set of vectors: r, {agg .a......m)) Each a-vector repre- sents an [S| -immensional hyper-plane, and defines the value function over bounded region of the belief:
insufficient. Consequently,
3. Asymmetric Constrained Partially Observable Markov Decision Processes (ACPOMDPs)
(b)
=max Σ Ea(s)b(s). acT, BES
Lemma 3.2. Assume the state space S, action space A, and observation space Z are finite. Let [S]. |A), and [Z| represent the number of states, actions, and observations, respectively. Let [Tt-a] denote the size of the solution set for the value function V--((b) at time step L- 1. The minimal number of elements required to express the value function Vz(b) at time step t, denoted as [Te), grows as [Tyl = O([A[[TT11)))) (Pineau et al., 2006).
In this section, we introduce our formulation of the Asym- metric Constrained Partially Observable Markov Decision Processes (ACPOMDPs) and compare it with the standard CPOMDPs. This comparison highlights the advantages of utilizing an asymmetric architecture.
3.1. Framework Setup
We propose ACPOMDPs, a relaxed variant of CPOMDPs. The key distinction is that ACPOMDPs assumes the avail- ability of the underlying states when computing the long- term expected values. Thus, ACPOMDPs are formulated by a tuple (5, B, A, T, RB, C, d, 7), and aims to maximize the
We conclude that, at each time step I, the belief state space can be represented as a discrete representation space that exactly captures the value function (b))). The size of this space is given by [T] = O([A) Z1). Furthermore, as derived from (6), in the ACPOMDPs framework, the
3

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
Exploration & Deployment
Stage :: World Model Learning
Stage Actor Critic Learning
Actor
Privileged Information
Privileged Representation
Privileged Critics
align
D
O
ITp
Object pose
Sensors
History actions
anc
Simulator
6.. Fz
Gesaeferia
0
for, tess a.....
for ac...
Naive WM
Privilaged
42
a2
%
Privileged Imagination
Naive WM
Actor
Privileged Predictors
Figure Overview of PIGDreamer. components highlighted purple are trained using privileged information, which collected from the simulators during exploration. We enhance the estimation of the critics by granting access to privileged information, thereby improving the actor's policy. Additionally, privileged representations are constructed provide the actor with more information.
required size of the representation space is reduced to [S]-
4.1. Overview
Clearly, |S| < [Tl = O[[A[]r;-]11))
We base our model on DreamerV3 (Hafner et al., 2023), renowned algorithm noted for its cross-domain general- ity and insensitivity to hyperparameters, wherein the world models are implemented
Thus, ACPOMDPs significantly reduce the size the repre- sentation space required to express the value function (b)), eliminating abservation-relater uncertainties the greatest extent possible. This, in turn, reduces the number of updates required for the critic to estimate the value function V (b).
Recurrent State-Space Model (RSSM) (Hafner al., 2019). In our framework, in- corporate additional components trained using privileged information to enhance policy training, subsequently dis- ubling these components during deployment to ensure that the policy operates solely on partial observations. Specifi- cally, during training, we concurrently develop two world models that serve as environment simulators: the naive world model and the privileged world model. The naive world model learns from the agent's observations, while the privileged world model learns from the underlying state. Subsequently, the actor and the critics are trained using the abstract sequences generated by these world models, where the actor operates exclusively on partial observations, while the critics are granted access to privileged information. Fi- nally, during deployment, only the naive world model and the actor, which operates partial observations, acti- vated. Figure I provides a clear visualization of our training pipeline. In summary, we leverage privileged information in three distinct ways:
Theorem 3.3. Let V* Vasym (b) and Vigm (b) represent the op- timal long-term expected values under the ACPOMDPs and CPOMDPs frameworks, respectively. Then, for all belief states E B, the inequality holds: Väsym (b) 2 Vegm(b) (The proof is provided in Appendix A)
The conclusion indicates that ACPOMDPs, by leveraging asymmetric information, yields superior policies compared to CPOMDPs. This is because, for any belief state e B, the optimal long-term expected reward under ACPOMDPs is always greater than equal to that under CPOMDPs. Regarding safety, ACPOMDPs provide more accurate es- timations due to additional information, while long-term expected costs under CPOMDPs are consistently lower or equal to those of ACPOMDPs. This implies that CPOMDPs tend to underestimate future safety risks.
Privileged Representations. Intuitively, providing the actor with a more informative representation results in superior policy. Inspired by this concept, enhance the representa- tions within the naive world model by aligning them with those derived from privileged information.
4. Methodology
Grounded in the ACPOMDP framework, we now present PIGDreamer, a Safe RL approach that utilizes privileged information to improve the policy of an agent operating under partial observability.
Privileged Predictors. Accurate prediction of rewards and

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforeement Learning
costs from partial observations is infeasible, as they do not
contain comprehensive information about the underlying states. We address this challenge by providing the predictors with privileged information.
Buaodas 1774 [e]
Dysanioo 1914 (i (4-3)
KL
Posterinnn 14:-5)
Based on our Theorem 3.3, granting critics access to privileged information can enhance policy performance. Consequently, we provide the critics with privileged information refine their estimations.
Privileged Critics.
Osacks Postetirn o, (o
Dysanies 2]44 (4--)---)))
Posittrr d/ 1n/-4/1
Poedeee 220h wh.46
Predtteee (l-47) Divrgoon
4.2. Privileged Information Guided World Model
The world models are parameterized by learnable network parameters Φ. Each world model operates on its model state At each time step i, the world models receive observation Ot, an action ar, and privileged information it as inputs. Encoders map or and i to embeddings & and =tt respectively. The dynamics predict the next states 8 based on the action ag. The posteriors utilize embeddings zi and zt and the predicted states 81+1, to predict the true underlying states 8t+1- Finally, reward and cost predictors use the concatenation of S and S to predict rewards and costs, while decoders employ their corresponding of and st to reconstruct the inputs and i, Detailed definitions of these components are provided below:
Figure 2. Visualization of the world model learning pipeline. Blue represents the components belonging to the naive world model, while purple indicates the components associated with the privi- leged world model.
Representation Alignment. Specially, in the naive world model, at each time step i, the oracle posterior predicts the oracle representations st, which encapsulate both the observation or and the privileged information i, We enforce the aligmment between the states s and the oracle states s* with the align loss defined in Equation 11 to enhance the state representations of the naive world model as follows:
Naive World Model
Latign = (s $t
(11)
Encoder: Po Pz 0z)
Privileged Predictors. In the PIG World Model, the predic- tors predict rewards and costs utilizing the most informative state representations s* and st. Here, s encompasses the observational information upon which the agent operates, while s+ contains the privileged information that enhances 8 prediction capabilities. The predictors are optimized with the following loss:
Decoder: Po( it 8t)
Dynamics: (8, 81-1: ac-1)
Posterior: (s 18, 2t
Oracle Posterior: 2o (8) (8) 87 22 ,
Privileged World Model
Encoder: P=(=;" iz)
Spred = In p (f+ ((,
In In (fre (8,888)
(12)
Decoder: (0 st))
Dynamies: Po(s) 15-1;
reward loss
cust loss
Posterior: ga(s; xt)
Afterwards, the decoders are supervised to reconstruct their corresponding inputs using the reconstruction loss in Equa- tion (13). Specifically, the decoder in the naive world model reconstructs the inputs o and i, using the oracle representa- tions 7f to capture all relevant information. This information is then distilled into 82 Compared to previous works (Lam- brechts et al., 2024; Hu et al., 2024) that directly reconstruct privileged information i, from the state representations s, this method significantly more robust when the privileged information is excessively informative for reconstruction:
Reward Predictor: po (f 87 s7)
Cost Predictor: (c s7,
The world models are classified into two categories: the naive world model and the privileged world model. The naive world model, which relies solely on the observation o, as input, is available to agents. As depicted in the Figure 2, the world models predicts state transitions using the dynam- ics by minimizing the dynamic loss between the predicted states & and the true underlying states 8t. The dynamic loss presented in Equation
Ldec = - -n (ōr, si) In (de $x
(13)
Ldyn ( s,) + Crep , ()) = KL lq || sg(p)] + KL (sg(4) g) |l pl. (10)
st),
(9)
observation
privileged information
Finally, the objective of the privileged information guided world model can be summarized as follows:
in the above, sg(-) represents the nop-gradient operator, while KT. denotes the Kullback-Leible (KI.) divergence.
L = Layn + Latign + Ldec + Lpred.
(14)
5

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforeement Learning
Actor
Privileged Crities
while the actor optimizes its policy according speci- fied objective. From the given imaginary trajectory, the bootstrapped (A() value R for the reward critic is calculated as follows:
(a .2)
Optimize
(T ft =ft+7(11-
(444))
(16)
Dysanio $395 (Sareea
+ +AR ,tt)), (
Racaal Proditae 22t (4-46)
rostder// (6.4)
[ (t) = Ve. (ST, s5)
(17)
Actenaai 4,)
Rovarl Collerpp60.1)
These values are used to assess the long term expected reward. where Vo. (s;) is approximated by the reward critic to consider the returns that beyond the imagination horizon H. Note that, we show here only the calculation of R^ (st), the calculation of (st is equivalent (16).
Con (65.e
Figure 3. Visualization of the actor critic learning pipeline. The lower part presents the generation process of the abstract sequences, dubbed Twisted Imagination. The upper part demonstrates the learning mechanism of the actor-critic model.
Lagrangian Constrained Policy Optimization. With the calculated TD(A) values (>,)) and (*(s)), we follow Equation (27) to define the policy optimization objective:
4.3. Privileged Safe Actor Critic Learning
The actor and critics learn purely from the abstract se- quences generated by world models. Specifically, at time step the actor, parameterized with the learnable network parameter 0, operates on the state 8f to predict the policy distribution (a | st ). The critics, on the other hand, operate on the state 8, and s!" to estimate the long-term expected returns (x, (*()) and (5, , Figure is presented for a clear visualization of our training pipeline. In summary, the key components of the actor and crities are:
[()) Σ sg (R + n/H [No (at [5]]
V (C) penalty term
(18)
This policy optimization objective encourages the actor to maximize the expected reward while simultaneously satis- fying the specified safety constraints. The penalty term is formulated using the Augmented Lagrangian method (Dai & Zhang, 2021). which penalizes behaviors that violate safety constraints. Additionally, an entropy term is included in the objective to promote exploration. Further details regard- ing the policy optimization objective and the Augmented Lagrangian method can be found in Appendix C.
Actor:
at (at 8
Reward Critic:
Uyt (sT , st) ) Re E. [RA] Uee. (85 , at) Rs E.g [CA]
(15)
Cost Critic:
Twisted Imagination (TI). Twisted Imagination is pro- cedure for generating abstract sequences from the world model, leveraging low-dimensional privileged information. As illustrated in Figure 3, starting from the representations of replayed inputs 5; and s, the actor samples an action at from the policy distribution (a s, ). utilizing s from the naive world model. Subsequently, each world model predicts its next representations $741 and st The cost è and reward f; can then be predicted by the respective pre- dictors until the time step t reaches the imagination horizon H = 15. This synchronization of the imagination process across the two world models enables the actor and critics to learn from coherent simulated trajectories.
5. Experiments
Experimental Setup. In all our experiments, the agent observes a 64x64 pixel RGB image from the onboard cam- era. We assess the task objective performance and safety using the following metrics:
Average undiscounted episodic return for E episode: ( = Σ E Tr.
Average undiscounted episodic cost return for E Tap
episode: J.(T) = E E C.
Privileged Critics After the Twisted Imagination, an ab- stract trajectory { st ae. ft, is provided to the actor and the critics. Similar to privileged predictors, the critics are supplied with privileged information s,t and the state s, to generate more accurate estimations. Specifically, based on the abstract trajectory, the critics can estimate the
We compute ((z) and Je((T) by averaging the sum of costs and rewards across E = 10 evaluation episodes of length Tep 1000. The results for all methods are recorded once the agent reached 4M environment steps. Detailed designs of privileged information, descriptions of all baselines, and additional experiments can be found in Appendix D.
long-term expected returns Ug (55 , sx) and Ug. (8, 8 )
6

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
MEDIAN
IQM
MEAN
Safe
20
LAMBDA
Eiformad--Deammerl.aa
SafeDreamer
Scaffolder(Lag)
920%
PIGDreamer(Ours
3,38
X78%
10.0
12.5
15.0
12
12
14
1.00
Human Normalized Scores
MEDIAN
IQM
MEAN
Safe-SLAC
11.07
LAMBDA
Informed-DeamerrLgg)
SafeDreamer
440
Scaffolden(Lag)
4:67
6.59
PIGDreamer(Ours)
2.5
5.0
7.5
4
2.5
5.0
7.5
Human Normalized Scores
ScatrAbe-
Safety Gymnasium benchmark results.
(b) Guard benchmark results.
Figure 4. Aggregate metrics for summarizing benchmark performance. In Figure 4a, we employ the rliable library (Agarwal et al., 2021) to compute the median, inter-quartile mean (IQM), mean estimates for normalized reward and cost returns, accompanied by 95% bootstrap confidence intervals (CIs) based on three rums across five tasks. In Figure 4b, we calculate the average rewards and costs across three tasks to compare the performance of various privileged variants SafeDreamer.
Results. As illustrated in Figure 4, our algorithm achieves state-of-the-art performance by leveraging privileged infor- mation. terms of safety, PIGDreumer not only meets the safety constraints at convergence, but also attains near- zero-cost performance. Concerning rewards, PIGDreamer outperforms all methods, especially its unprivileged variant, SafeDreamer which demonstrates significant improvements. Conversely, as depicted in Figure 6, LAMBDA matches PIGDreamer in rewards but fails to ensure safety and does not produce a practical policy for the PointButton/ task. Ad- ditionally, ALAc performs worse, failing in the Point- Pushl and RacecarGoall tasks and exhibiting considerable safety constraint violations. As illustrated in Figure 4a, Scaf- folder (Lag) and Informed-Dreamte (Lag) exhibit superior performance compared to SafeDreamer. This finding indi- cates that the utilization of privileged information enhances the agent's 's capacity to Iearn a more effective policy, even with limited observations, which aligns with Theorem 3.3. However, as depicted Figure 4b, these algorithms expe- rience decline in performance relative to SafeDreamer after incorporating privileged information. This paradox arises because, the agent is unable to compensate for the information gap pertaining to the privileged information, resulting in sub-optimal policies. In contrast to alternative methods, PIGDreamer consistently demonstrates superior performance. This underscores the robustness of our ap- proach in utilizing privileged information.
469%
Training Time (h)
443%
25.25
20
21.46
288%
A10%
19.14
16.45
10
0
Distil(Lac)
SceffoldeniLagl PiGD (Ouns)
SafeDreemer
Informed-Dreametn (Lag)
Figure 5. Training Efficiency Comparison: We visualize the aver- age training time for the tasks on Guard henchmark.
exploiting privileged information yield disparate results. Distill (Lag), which trains a privileged teacher policy and distills the knowledge to the student policy, suffers from severe performance degradations resulting from the infor- mation gap between the privileged information and partial observations. Informed-Dreamme reconstructing privileged information from partial observa- tions, achieves only a marginal performance improvement. In contrast, Scaffolder builds on Informed-Dreamer by pro- viding predictors and critics with privileged access and in- corporating an additional privileged actor for exploration, leading significantly improved outcomes. PIGDreamer consistently achieves sfate-of-the-an art performance across all benchmarks. In particular, in the Guard benchmark, PIG- Dreamer outperforms alternative variants with a 136% per- formance improvement. Nonetheless, these enhancements result in extended training times. Our analysis in Figure.5 demonstrates that Scaffolder requires an average training
which relies solely on
Benchmarking Privileged Variants. We now examine the privileged variants: Scaffolder (Lag), (Lag), Distill (Lag), and PIGDreamer. The comparisons presented Figure reveal that different approaches to

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
CarGoall
PotzsButtol
PotutCioal
PoinsPushl
RacecarCioal
105
165
145
14h
116
Aocree con.ec..ooooo ream
13-
fotrogestt cooegensteou
20-
10-
Sea))) PIGrreamet(Uus) Informod-
Safe-SAC
LAMBDA
SafeDreamer
Figure 6. The experimental results for the Safety-Gymnasee Gymnasium Benchmark. The upper figures illustrate the curves of episodic reward. while the lower figures depict the curves of episodic cost. The red solid line indicates the cost constraint for this task.
time that is 69% longer compared SafaDremmer, while Informed-Dreamer necessitates only 10% longer average training time. Our algorithm, PIGDreamer, achieves the best outcomes with only a 28% increase in average training time. This indicates that our algorithm exhibits the highest efficiency in utilizing privileged information.
on more informative representations.
Task
Method
Reward
Cost
NLI
10.79
0.41
SafetyPointGoal2
TI (Ours)
13.59
0.73
NLI
14.79
0.64
SafetyCarGoall
TI (Ours)
17.32
0.43
IQM
NLI
13.99
1.57
SafetyRacecarGoall
PIG -No Rep PIG -Unprivileged PIGDreamer(Ours)
TI (Ours)
11.38
0.83
Table I. Comparison with Nested Latent Imagination (NLI).
13.5 Human Normalized Scores
10.5
12.0
Table I presents comparison of various trajectory gen- eration procedures. The results indicate that TI achieves competitive performance compared to NLI, while utilizing lighter model design, which is significantly important for robustness and versatility.
Figure 7. The ablation study results.
Ablation Study. Now assess the component-wit wise con- tributions of the privileged components in our algorithm. Our ablation study includes the following settings: (1) PIG -No Rep: In this variant, the naive world model does not align the state representations with the oracle representa- tion. (2) PIG Unprivileged: In this variant, no privileged information is accessible during training. (3) NLI: In this variant, replace Twisted Imagination with Nested Latent Imagination (NLI). Figure.7 presents the results of the abla- tion study conducted on the Safety-Gyy Gyynnssimm Benchmark, while detailed experimental results can be referenced in Fig- ure.9. As illustrated in Figure.7, PIG No Rep, shows only marginal improvement compared to PIG Unprivileged. This limited enhancement occurs because the privileged critic mainly improves the agent's policy by accurately esti- mating the value function, which guides decision-making but is insufficient for substantial improvements. Conversely, PIGDreamer achieves significant improvements by utilizing privileged representations. This because the privileged representations enable the agent to predict privileged infor- mation from its observations, allowing the agent to operate
6. Conclusion
We propose PIGDreamer, a model-based RL approach specifically proposed for partially observable environments with safety constraints. This approach formalized within the ACPOMDPs framework, providing theoretical advan- tages in addressing the challenges of partial observability and safety. By leveraging privileged information through privileged representation alignment and an asymmetric actor-critic structure, PIGDreamer achieves competitive per- formance using vision-only input on the diverse tasks. Our experiments show that PIGDreamer is the most effective approach for utilizing privileged information within world models, excelling in performance, safety and efficiency. We believe PIGDreamer marks significant advancement in harnessing RI for real-world applications. However, we noted that privileged information did not always lead to im- provements in certain tasks, necessitating further research to explore the relationships between specific types of privi- leged information and particular tasks.
8

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
Acknowledgement
Cai, Y., Liu, X., Oikonomou, A., and Zhang, K. Prov- able partially observable reinforcement learning with privileged information, 2025, URL httpt:tttaaaaee.. org/ abs /2412.4222 00985.
This work was supported by the National Natural Science Foundation of China (Grant No. 62272024).
Impact Statement
Dai, Y.-H. and Zhang, L. The augmented lagrangian method approximately solve convex optimization with least constraint violation, 2021. URL https:
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must specifically highlighted here.
/ //arxiiv.orggeeee. ht///bb//2111 wtt.:///abb//1..111.......9 06199..
Egorov, M., Sunberg, Z. N., Balaban, E., Wheeler, A., Gupta, J. K., and Kochenderfer, M. J. Pomdps.jl: A framework for sequential decision making under uncer- tainty. Journal of Machine Learning Research, 18(26):1- 5, 2017. URL 1 16-300 .ttml.
References
Achiam, and Amodei, D. Benchmarking safe ex-
ploration in deep reinforcement learning- 2019.
URL https://apl http://app..senniisssssttte semanticscholar.co org/ CoxpusID::28283322C.
Feng, S., Sun, H., Yan, X., Zhu, H., Zou, Z., Shen, S., and Liu, H. X. Dense reinforcement learning for safety validation of autonomous vehicles. Nature, 615(7953): 620-627, 2023.
Achiam, Held, D., Tamar, A., and Abbeel, P. Con- strained policy optimization, 2017. URL https:// arxiv. www.or///bo.... /105.
Hafner, D., Lillicrap, T., Fischer, I,, Villegas, R., Ha, D., Lee, H., and Davidson, J. Learning latent dynamics for planning from pixels, 2019. URL https:: //ariii.//rxv...
Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A., and Bellemare, M. G. Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems, 2021.
org/ abs////////// 1811.11111 04551.
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination, 2020. URL wwwpthtp///// a www.i.vvv.... www.v.ww.//..wwwow........w.ooow /abs////////// / 91 2. 01 603.
Altman, E. Constrained Markov decision processes. Rout- ledge, 2021.
Hafner, D., Pasukonis. J., Ba, J., and Lillicrap. T. Mastering diverse domains through world models. arXiv preprint 2rX1v:2301.04104 2023.
As, Y., Usmanova, I., Curi, S., and Krause, A. Constrained policy optimization via bayesian world models, 2022. URL
org/ /wt.///11.... 09802.
Hogewind, Y., Simao, T. D., Kachman, T., and Jansen, N. Safe reinforcement learning from pixels using a stochastic latent representation, 2022. URL https://arxttti org abs / 2200. 01801.
Avalos, R., Delgrange. F., Nowe, A., Perez, G., and Roijers, D. M. The wasserstein believer: Learning belief updates for partially observable environments through reliable latent space models. The Twelfih International Confer- ence on Learning Representations, 2024. URL https: //openrevies net. forum? Let td-.er
Hu, E. S., Springer, J., Rybkin, O., and Jayaraman, D. Priv- ileged sensing scaffolds reinforcement learning, 2024. URL httph:tp////mmm arxiv.ooorr.oo...o..aa....a.aaae /as//ab//// /tt.//405.15555
Baisero, A. and Amato, C. Unbiased asymmetric reinforce- ment learning under partial observability. arXiv preprint arXiv:2105.11674, 2021.
Huang, W., Ji, J., Xia, C., Zhang, B., and Yang. Y. Safe- dreamer: Safe reinforcement learning with world mod- els, 2024. URL https: /tapp::///axxv...// org/ abs/23077 07176.
Baisero, A., Daley, B., and Amato, C. Asymmetric dqn for partially observable reinforcement learning. In Conference on Uncertainty in Artificial Intelligence, 2022. URL https:/ http://pp...mmanicccchcooooooooo org/Corpusthp http:/oopuSTD::288995500 252898960
Jayant, A. K. and Bhatnagar, S. Model-based safe deep reinforcement learning via a constrained proximal policy optimization algorithm, 2022. URL https: //arxiv. org/ abs 2210.2211... 07573.
Berkenkamp. F., Turchetta, M.. Schoellig, A. P., and Krause, A. Safe mrodel-based reinforcement learning with stability guarantees, 2017. URL https: :ttps://arraae /arxiv...v./////////aaa///aaavva org/ abs/ 7705.08557
Ji, J., Zhang, B., Zhou, J., Pan, X., Huang, W., Sun, R., Geng, Y., Zhong, Y., Dai, J., and Yang, Y. Safety- gymnasium: A unified safe reinforcement learning bench- mark. arXiv preprint arXiv:2310. 1266,, 2023a.
9

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
Ji, J., Zhang, B., Zhou, J., Pan, X., Huang, W., Sun, R., Geng, Y., Zhong, Y., Dai, J., and Yang. Y. Safery- gymnasium: A unified safe reinforcement learning bench- mark, 2023b. URL https: / arxiv. org/ abs/ 2310. 1...12565 12567.
Li, J., Fridovich-Keil, D., Sojoudi, S., and Tomlin, C. J. Augmented lagrangian method for instantaneously con- strained reinforcement learning problems. In 2021 60th IEEE Conference on Decision and Control (CDC), pp. 2982-2989. IEEE, 2021.
Ji, J., Qiu, T., Chen, B., Zhang, B., Lou, H., Wang, K., Duan, Y., He, Z., Zhou, J., Zhang. Z., Zeng, F., Ng, K. Y., Dai, J.. Pan, X., O'Gara, A., Lei, Y., Xu, H., Tse, B., Fu, J., McAleer, S., Yang, Y., Wang, Y., Zhu, S.-C., Guo, Y., and Gao, W. Ai alignment: comprehensive survey, 2024. URL https: /ate:://arre http///aaxvvv...aae/e// htrgg/abs/2310.. 19852.
Liu, Y., Halev, A., and Liu, X. Policy learning with con- straints in model-free reinforcement learning: A survey, In The 30th international joint conference on artificial intelligence (ijcai), 2021.
Moerland, T. M., Broekens, J., Plaat, A., and Jonker, C. M. Model-based reinforcement learning: A survey, 2022. URL https: /ttp:://aaxxvvoo/
Jin, C., Kakade, S. M., Krishnamurthy, A., and Liu, Q- Sample-efficient reinforcement learning of undercom- plete pomdps, 2020. URL httts://arxivv
abbs22006.67.
Nocedal, J. and Wright, S. Numerical optimization, pp- 1-664. Springer Series in Operations Research and Fi- nancial Engineering. Springer Nature, 2006.
abs/2006. 12484.
Kaelbling, I.. P., Littman, M. L., and Cassandra, A. R. Plan- ning and acting in partially observable stochastic domains. Technical report, USA, 1996.
OpenAI, Andrychowicz, M., Baker, B., Chociej, M., Joze- fowicz, R., McGrew, B., Pachocki, J., Petron, A., Plap- pert, M., Powell, G., Ray, A., Schneider, J., Sidor, S., Tobin, J., Welinder, P., Weng, L., and Zaremba, W. Learning dexterous in-hand manipulation, 2019. URL https: /apps://arxiin wttpaaxivvvvv.....a org. abs 808 00177.
Kaelbling, L. P., Littman, M. L., and Cassandra, A. R. Plan- ning and acting in partially observable stochastic domains. Artificial intelligence, (111)):99-113) 1998.
Krishnamurthy, A., Agarwal, A., and Langford. J. Pac reinforcement learning with rich observations, 2016. URL https: / /arxivvvvvv org/ otg/abs/////oooo.o...........err / 102...6222 02722.
Papadimitriou,, C. H. and Tsitsiklis, N. The complexity of markov decision processes. Mathematics operations research, 12(3):441-400, 1987.
Lambrechts, G., Bolland, A., and Ernst, D. Informed pomdp: Leveraging additional information in model- based rl, 2024. URL htap::://axxiv..oo /
Pineau, J., Gordon, G., and Thrun, S. Anytime point- based approximations for large pomdps. Journal of Ar- tificial Intelligence Research, 27:335-380, November 2006. ISSN 1076-9757. doi: 10.1613/jair.077.. URL http: /tp///dxxooo doi org.org//////// /tl.g//...11... 11.1100..13//// 111r//10.6133/aar.2222
2306. 11488.
Langley, P. Crafting papers on machine learning. Langley, P. (ed.), Proceedings of the 17th International Conference Machine Learning (ICMI. 2000), pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
Pinto, L., Andrychowicz, M., Welinder, P., Zaremba, W., and Abbeel, P. Asymmetric actor critic for image-based robot learning, 2017. URL https: htttps://arxvv.ogg abs 0tt./1110..6655
LeCun, Y. and Courant. wards autonomous machine intelligence. http://hhh..cmmmtcccccaaoo//Cond 2022.
A path to-
Polydoros, A. S. and Nalpantidis, L. Survey of model- based reinforcement learning: Applications on robotics. Journal of Intelligent & Robotic Systems, 86(2):153-173, 2017.
Lee, A. X., Nagabandi, A., Abbeel, P., and Levine, S. Stochastic latent actor-critie: Deep reinforcement learn- ing with latent variable model, 2020. URL https: //arxiv. erg/abs/ 907. 00953.
Salter, S., Rao, D., Wulfmeier, M., Hadsell, R., and Posner, I. Attention- reinforcement learning, 2021, URL https: https://axxxvvvvie abs 1911 08363.
Lee, J., Kim, G.-H., Poupart, P., and Kim, K.-E. Monte- carlo tree search for constrained pomdps. Advances in Neural Information Processing Systems, 31, 2018.
Thomas, G., Luo, Y., and Ma, Safe reinforcement
learning by imagining the near future, 2022. URL https: httpa/axvvv.g////e htt.b////222.07777.
Li, J., Koyamada, S., Ye, Q., Liu, G., Wang, C.. Yang, R., Zhao, L., Qin, T., Liu, T.-Y,, and Hon, H.-W. Suphx: Mas- tering mahjong with deep reinforcement learning, 2020.
Yamada, J., Rigter, M., Collins, J., and Posner, I. Twist: Teacher-student world model distillation for efficient sim- to-real transfer, 2023. URL https: / //t/x/ivvvco org/ abs/2311 033.222.1.....2...
URL http///awwvv.reeaeeeeee wwwpobrgg//.oa.a// 13590.
10

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
Zanger, M. A., Daaboul, K., and Zöllner, M. Safe contin- uous control with constrained model-based policy opti- mization. In 2021 IEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pp. 3512-3519. IEEE, 2021.
Zhang. Y., Vuong, Q- and Ross, K. W. First order constrained optimization in policy space, 2020. URL /taps://arxivv org/ abs//b///////// 2002 06506.
Zhao, W., Chen, R., Sun, Y., Liu, R., Wei, T., Li, F., and Liu, C. GUARD: A safe reinforcement learning benchmark, 2024. URL https: / http//ppweeee.eeeeeeteeee wewpeeewwwwwweeeeeeeeeeeeeeeeeee forum? id=LkIvfyyh6v.

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforeement Learning
A. Proof
In this section, we prove Theorem 3.3,
First, we donate the optimal long-term values in belief space under the
ACPOMDPs framework as Vas Vasym (6):
V asym (b)
Σ b(s)
#ES
(19)
R(s, + Ep(s)s
(*(s)
max a
where (s) represents the optimal long-term values in the state space. Similar to V Vasy (b), the optimal long-term values in the belief space under the CPOMDPs framework are represented as V (b):
V sym (6) = max aEd R(b,a) + Σ
Pr(z|b, sym (b4.2)
(20)
ZEZ
Since
Pr(z(b,a)= = Σ Pr(z]a, 8') Σ
Pr(s'(s,a)b(s)
(21)
s'ES
sES
sym (b) can be rewritted as:
V sym
(6) max
R(b, a) + Σ
(b4.3)
Σ
Pr(zla,s') Σ
P(5'(s,a)((s)
sym
atA
zEZ
s'ES
sE8
(22)
max .εΛ
R(b, a) + Σ
Ve
sym
(baz)
ΣΣ
Pr(z|a, s (P(ss)s,a)b(s))
ZEZ
sES
Proof. According to Lemma 3.1, (s') can be expressed by set of vectors: r, {ao, an am As result, V*('') can be rewrite as the following equation:
(s) = max
R(s,a) + Σ
P(5'(s, max a(s')
(23)
aEA
aeT
Similarly, Viym (b) can be rewritten as:
V sym (b)
max aEA
Σ R(s,)))(s))+
+ Σ max Σ Σ
(24)
aeT t-1
LIES
#EZ
sE8 eES
Then:
7 Vasy (b) Σ V(s)V*(s)
= Σ
B(8) max DEA
R(8,a)+ + Σ P(s)(s,a) max ael.
a(s')
SES
2 max a
(ssR((s))) +7 Σ b(s)) Σ
P(s'ls, a) max a(s')
aeF:
LEES
sES
F'S
(25)
R(b, +7 ΣΣ Σ
Pr(zla, (()(sl,
maxa(s')
max A
det.
2
2 max asA
R(b,a) +7 Σ
max
ΣΣ
Pr(zla, P((el
aeT, sEz
#ES
(6)
sym
12

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
B. Hyperparameterss
B.1. PIGDreamer & SafeDreamer Informed-Dreamer & Scaffolder
Table 2. Hyperparameters
Name
Symbol
Value
World Model
Number of latent classes
48
Classes per latent
48
Batch size
B
64
Batch length
T
16
Learning rate
10-4
Coefficient of KL divergence in loss
Oigy Op Bo, Pr. Pe
0.1, 0.5
Coefficient of decoder in loss
1.0, 1.0, 1.0
Planner
Planning horizon
H
15
Number of samples Mixture coefficient
Nn
500
M
0.05
N M N
Number of iterations
J
6
Initial variance
PID Lagrangian
Proportional coefficient
Kp
0.01
Integral coefficient
K2
0.1
Differential coefficient
Ka
0.01
Initial Lagrangian multiplier
Xzo
0.0
Lagrangian upper bound Maximum of X
0.75
Augmented Lagrangian
Penalty term
v
5-9
Initial penalty multiplier
1-6
Initial Lagrangian multiplier
X
0.01
Actor Critic
Sequence generation horizon
15
Discount horizon
?
0.997
Reward lambda
A
0.95
Cost lambda
A
0.95
Learning rate
3 10
General
Number of other MLP layers
5
Number of other MILP layer units
512
Train ratio
512
Action repeat
4
B.2. Safe-SLAC
Hyperparameters for Safe-SLAC, We maintain the original hyperparameters unchanged, with the exception the action repeat, which we adjust from its initial value of 2 to 4.
13

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
Table 3. Hyperparameters for Safe-SLAC
Name
Value
Length of sequences sampled from replay buffer
15
Discount factor
0.99
Cost discount factor
0.995
Replay buffer size
2 x 10 32
Latent model update batch size Actor-critic update batch size Latent model learning rate Actor-critic learning rate
64
1 x 10
X 10
Safety Lagrange multiplier learning rate Action repeat Cost limit
x 10
2.0
Initial value for Initial value for À Warmup environment steps Warmup latent model training steps
4 x 10-3
2 x 10
60 x 103
30 x 103
Gradient clipping max norm
40
Target network update exponential factor
10-3
B.3. LAMBDA
Hyperparameters for LAMBDA. We maintain the original hyperparameters unchanged, with the exception of the action repeat, which we adjust from its initial value of 2 to 4.
Table 4. Hyperparameters for LAMBDA
Name
Value
Sequence generation horizon
15
Sequence length
50
Learning rate
1 10-1
Burn-in steps
500
Period steps
200
Models
20
Decay
0.8
Cyclic LR factor
5.0
Posterior samples
5
Safety critic learning rate
2 10
Initial penalty
5 10
Initial Lagrangian
1 10-6
Penalty power factor
1 x 10
Safety discount factor
0.995
Update steps
100
Critie learning rate
8 10-5
Policy learning rate
8 10
Action repeat
4
Discount factor
0.99
TD(A) factor
0.95
Cost limit
2.0
Batch size
32
14

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforeement Learning
C. The Augmented Lagrangian
The Augmented Lagrangian method incorporates the safety constraints into the optimization process by adding a penalty term to the objective function. This allows the actor model to optimize the expected reward while simultaneously satisfying the specified safety constraints. As result, by adopting the Augmented Lagrangian method, we transform the optimization problem in (4) into unconstrained optimization problem:
C (
- b,) max min R(T) Σ -() E X
1
Σ
x
(26)
==1
where X are the Lagrange multipliers, each corresponding to a safety constraint measured by Ci(T), and jk is a non- decreasing penalty term corresponding to gradient step k. We take gradient steps of the following unconstrained objective:
R(T; uR) R(TT)= Σ ((C,(T),
(27)
We define A, = Ci(T) be The update rules for the penalty term (TC((T)) Xkx ukx) and the Lagrange multipliers X are as follow:
X A + 2
A + + +
Vi [m] (L (2),
(At)
(28)
otherwise.
2pt
D. Experiments
D.1. Privileged Information Design
In different tasks, it is necessary to customize the use of different privileged information, and different privileged information will have different impacts, we show our privileged information settings in our experiments.
Privileged Information Name
Dimension
Deseription
bazards
[n, 3)
Represents the relative positions of hazards the eavironment. containing 3D coordinates [2. a).
velocimeter
[3.]
Provides the agent's 's velocity information in three-dimensional space [es, My. Val-
accelerometer
(3.
Provides the agent's 's acceloration information in three-dimensional space (g:
gyro
13.
Provides the ageut's is angular velocity information (ola Wy 60g]
goal
(3,
Represents relative conrdinates of the target position the agent needs reach ("goal: Mgual- Egeal:
robot.M
3)
Represents rotation matrix robot, describing the robot's ocientation rotation dinee-dimensioual space.
1 I action
[4.]
Represents the action information from previous time step.
past actioo
(1.
Represents the action information from the second-to- time step.
past action
[4.]
Represents the action information from the thard-to-last time step.
culer
12.
Represents the agent's 's pose information given in Euler angles poll,,
Table 5. Privileged Information: SafetyQuadrotorGoall
D.2. Experimental Setup
Setup. Our experiments were conducted using the following configuration: single A A100-PCTE-40GB GPU (40GB), 10 vCPU Intel Xeon Processor (Skylake, IBRS), and 72GB of memory.
Baselines. We compared PIGDreamer to several competitive baselines to demonstrate the superior results of using privileged information. The basclines include: 1. (Lag)) (Hu et al., 2024) Integrates Scaffolder with the Lagrangian methods. (Lambrechts et al., 2024) Integrates Informed-Dreame SafeDreamer: (Huang et al., 2024) Integrates Dreamerv3 with the Lagrangian methods. 4. LAMBDA: (As et al., 2022) novel model-based approach utilizes Bayesian world models and the Lagrangian methods. 5. Safe-SLAC: (Hogewind al., 2022) Integrates SLAC with the Lagrangian methods. CPO: Achiam al. (2017) the first general-purpose policy search algorithm for constrained reinforcement learning with guarantees for near-constraint satisfaction each iteration. 7. PPO Lag: Achiam & Amodei (2019) Integrates PPO with the Lagrangian methods. 8. TRPO Lag: Integrates TRPO with the Lagrangian methods. 9. FOCOPS: Zhang et al. (2020) initially determines the optimal update policy by addressing a constrained optimization problem within the nonparameterized policy space. Subsequently. FOCOPS projects the update policy back into the parametric policy space. Notably, the OSRP, OSRP Lag and SafeDreamer are three algorithms proposed by SafeDre
2. Informed--Deeamr
with the Lagrangian methods.
15

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
D.3. Addition Experiments
Gozzart
GoalAnttGhosts
GoalHumanoid&Hazards
Average
Methods
Reward
Cost
Time
Reward
Cost
Time
Reward
Cost
Time
Reward
Cost
Time
SafeDreamer
2.54
1.00
14.27
10.67
0.93
16.7
0,34
10.8
13.94
4.52
4.24
14.97
Distill Dreamer(Lag)
1.32
10.52
20.12
0.50
1.02
23.21
1.41
24.07
21.04
1.08
11.87
21.46
Informed- Srammedd-Deamer(Lagg
1.32
0.03
17.03
10.09
2.54
18.28
-1,60
10.72
14.03
3.27
4.43
16.45
Scaffolder(Lag)
1.86
0.04
24.82
6.15
0.83
27.14
2.13
13.14
23.8
3.38
4.67
25.25
VIGDreamer(Oury)
14.18
0.92
18.45
15.76
2.28
20.86
2.09
10.56
18.12
10.68
4.59
19.14
Table 6. The experimental results for the Guard Benchmark.
Aoogg Nomgrenniiin
163
Aucrag- ntt
PIGDreamar((Ous) Privillged-Agemm
SafeDremmer
Figure Privileged Agent PointButton! Task.
Discussion Although Figure 8 demonstrates the superior performance of exploiting privileged information, it does not achieve significant improvements in certain tasks. To provide more detail regarding our experiments, Figure 6 presents the learning curves of all algorithms. Notably, we observe that PIGDreamer does not outperform SafeDreamer in the PointButton/ task. We conducted further experiments to investigate this phenomenon, and our findings indicate that even when the agent is granted access privileged information, still fails achieve higher performance. Consequently, this finding reveals that, in the PointButton| task, which involves moving objects in the environment, such privileged information does not yield any information gain.
D.4. model-free
We also compare our results with several model-free algorithms. In our experiments, the baseline algorithm results derived from SafePO Ji et al. (2023a), which is configured with cost limit of 25, whereas PIGDreamer is set with cost limit of 2. The comparison results are presented in the Table 7.
CPO
FOCOPS
PPO-Lag
TRPO-Lag
PIGDreamer
Tasks
Reward
Cost
Reward
Cost
Reward
Cost
Reward
Cost
Reward
Cost
CarGoall
23.2++1.9 1,9
28.2+4.66
21.5+000
28. ++000
13,8+3.3
13.4+10.88
22.2+3.9
26.2+66.1
14.5+0.5
4,2+
PointButton
6.8++166
29.8+6.1
8.9+
1..2+4.5
4,0+1.4
28.2+13.8
1.5+1.4
26.3++600
9.6+3.5
12.5+2.3
PointPush1
4.8.0.00
25.540.0
0.740.77
23.0421.1
0.6:1..3
26.2422.1
0.610.1
21.7+11.2
15.6:0.8
0.4:0.2
RacecarGoall
10.41:.2
29.427.0
4.5+2.2
93.7433..3
2.342.1
28.3
9.543.0
25.1+5.7
18.2+112
6.8+1.2
Average 11.3
28.3
8.9
38.8
5.2
26.6
9.9
24.9
14.5
6.0
Table 7. Comparison with model-free algorithms
As illustrated in Table 7, the model-free algorithm encounters difficulties in achieving higher rewards, even with a more relaxed cost threshold. This challenge arises partly from the inability of these algorithms to leverage historical information,
16

---

PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning
CarGoall
PoiniGoal2
PoiniPushl
RacecurGoull
Airenge scoand sctu
Avringe rowand sutum
15
10
10
le6
le6
leb
166
60
Avergg cont retuss
Asange cnd naturn
10
165
166
1c5
166
PIGDreamer(Ours)
PIG -No Rep
PIG -Unprivileged
Figure 9. Aggregate metrics for summarizing benchmark performance. We utilize the rliable library (Agarwal et al., 2021) to compute median, inter-quartile mean (IQM), and mean estimates for normalized reward cost returns, accompanied by 95% bootstrap confidence intervals (CIs) based on three runs across five tasks.
well as their lack of access to essential privileged information. In contrast, PIGDreamer significantly outperforms the baseline algorithm in both reward and safety, owing to its effective utilization of privileged information.
D.5. Training efficiency
Since model requires modeling world models, will compare the training efficiencies of the different algorithms to determine if the addition of world model results in significant increase in training time. Table 8 displays the hours
LAMBDA
Safc-SLAC
Informed Informed-Dreamer((Leg
Scaffoldcr(Lag)
SafeDreamer
Ours
Incremental
CarGoall
34.1
18.3
25.77
45.12
25.51
36.02
41.2%
PointButton!
28.4
31.4
31.61
51.02
29.42
34.59
17.6%
PointGoal2
31.1
22.5
31.51
50.68
29.52
34.52
16.9%
PointPushl
36.7
18.3
25.80
43.87
23.51
29.43
25.2%
RacecarGoall
34.1
20.7
31.08
49.54
29.59
33.84
14.3%
Average 33.2
17.7
29.16
48.05
27.51
33.68
22.4%
Table 8. Comparison of Training Time
required to train all algorithms for 2 million iterations each task. The comparison between SafeDreamer and PIGDreamer reveals that the inclusion of an additional world model in PIGDreamer resulted in an average increase of 22.4% in training hours.
SafetyPointGoal2
Saretya
SafetyPointPushl SafetyRacecarGoall
SafetyPointButton!
Methods
Reward
Cost
Reward
Cost
Reward
Cost
Reward
Cost
Reward
Cost
Distill
6.59
4.54
11.74
3.08
10.21
3.47
12.26
12.37
6.77
11.12
Ours
11.61
1.31
17.37
0.93
10.99
1.17
17.10
1.15
5.97
2.01
Table 9, Comparison of Safety Methods
17


Reviews:
[{"review_id": "sevr1Gj4db", "reviewer": "Reviewer_byT1", "summary": "Disclosure: I am an emergency reviewer, and also quite familiar with the paradigm of RL with privileged information.\n\nThe authors propose to exploit privileged information for policy learning in the context of safe RL. They use a world model approach like Dreamer, where an actor-critic is trained solely in a learned MDP (the world model). They use privileged information to train a privileged world model. Then, they exploit the privileged information in a variety of ways, such as aligning the privileged WM and unprivileged WM representations, using asymmetric critics and reward estimators, and generating imagined data with the privileged WM. They evaluate their method on the Safety Gymnasium benchmark, which is state-based goal-reaching tasks, and the agent must avoid hazards.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our gratitude to Reviewer byT1 for your insightful comments.\n\n**Q1. The mismatch arises from the Twisted Imagination (TI) and the ablation of TI with the Nested Latent Imagination (NLI).**\n\nThank you for your question. Indeed, the trajectories generated from TI exhibit variability due to their stochastic nature. To address this problem, NLI proposed in Scaffolder forces the matching between two trajectories by encoding the embedding from  $Path_2$ to $Path_1$. However, our ablation study about the NLI and TI, demonstrate that this method yields no performance improvement. We attribute this phenomenon to the robustness of the TI predictors. Although the representations of the privileged world model deteriorate during the imagination process, the predictors continue to make accurate predictions from the states $s^{*}$  and $s^{+}$ . \n\n| Task | Method | Reward | Cost |\n| --- | --- | --- | --- |\n| SafetyPointGoal2 | NLI | 10.79 | **0.41** |\n|  | **TI (Ours)** | **13.59** | 0.73 |\n| SafetyCarGoal1 | NLI | 14.79 | 0.64 |\n|  | **TI (Ours)** | **17.32** | **0.43** |\n| SafetyRacecarGoal1 | NLI | **13.99** | 1.57 |\n|  | **TI (Ours)** | 11.38 | **0.83** |\n\n**Q2. The implementation of Representation Alignment.**\n\nIn our work, we pull both state encoders towards each other  using the method proposed in Dreamerv3.\n\n**Q3. Clarification regarding the privileged representation.**\n\nIn our experiments, Informed-Dreamer exhibits instability in reward when the privileged information is too informative to be reconstructed from the history of partial observations. To address this issue, we enhance the state representation through **privileged representation alignment**. This approach enables us to establish correspondence between modalities (images and states), while ensuring that the training objective remains agnostic to the dimensionality of the privileged information.\n\n**Q4. Comparison with Scaffolder.**\n\nThank you for your inquiry. We are pleased to compare our algorithm with Scaffolder (ICLR24). We will evaluate our algorithm against Scaffolder (ICLR24) in the following aspects:\n\n1. **Theoretical Foundation:** We propose our method within the ACPOMDP framework while Scaffolder has made limited advancements in this area.\n2. **Privileged Representations:** see **Q3**.\n3. **Privileged Predictors:** The scaffolder comprises two groups of predictors: one group relies exclusively on the state of the naive world model, while the other depends on the state of the privileged world model. In our approach, we simplify these components by enabling the predictors to operate based on the states of all world models.\n4. **Privileged Imagination:** see **Q1**.\n5. **Privileged Exploration:** We eliminate this component for the following reasons:\n    1. It prolongs the training time.\n    2. The privileged actor operates using information that is inaccessible to the naive actor, potentially creating a significant disparity between their behaviors. As a result, the trajectories collected by the privileged actor become difficult for the naive actor to learn from.\n6. **Privileged Critics:** same with Scaffolder.\n7. **Efficiency:** Our method demonstrates enhanced efficiency in utilizing privileged information.\n\n**Q5. The teacher-student distillation baseline.**\n\nThank you for your valuable suggestion. We have added the teacher-student distillation baseline within the limited timeframe. Our experimental results demonstrate that our method consistently outperforms the distillation baseline by a significant margin. Furthermore, we observe that, in the absence of explicit safety constraint objectives, the student fails to meet safety constraints.\n\n| Task | Model | Training time | Cost | Reward |\n| --- | --- | --- | --- | --- |\n| SafetyPointGoal2 | Distill | 30.02 | 4.54 | 6.59 |\n|  | **Ours** | **29.71** | **1.31** | **11.61** |\n| SafetyCarGoal1 | Distill | 29.21 | 3.08 | 11.74 |\n|  | **Ours** | **28.87** | **0.93** | **17.37** |\n| SafetyRacecarGoal1 | Distill | **30.19** | 3.47 | 10.21 |\n|  | **Ours** | 31.03 | **1.17** | **10.99** |\n| SafetyPointPush1 | Distill | **31.03** | 12.37 | 12.26 |\n|  | **Ours** | 34.87 | **1.15** | **17.10** |\n| SafetyPointButton1 | Distill | 30.02 | 11.12 | **6.77** |\n|  | **Ours** | **29.54** | **2.01** | 5.97 |\n\n**Q6. Experimental setup.**\n\nThank you for your question. Our experiments actually use 64 × 64 pixel images for agent observations, while using low-dimensional state as privileged information. We will clarify this experimental setup in the final version of our manuscript. \n\n**Q7. Wasserstein Believer.**\n\nWe plan to incorporate a discussion of its relevance in the next version of our paper.\n\n**Q8. The stability of the augmented lagrangian method.**\n\nWe regret our inability to engage in further discussion on this matter due to space limitations. Given that the augmented Lagrangian method is a widely accepted practice within the Safe RL community, we have not investigated its stability.", "claims_and_evidence": "The overall system is verified by showing improvements over baselines. The improvement in returns and safety violation, compared to baselines, is marginal to moderate in my opinion, especially compared to the most competitive baseline Scaffolder. However, this method does seem more simple and faster to run which is good.\n\nThe design choice of certain components in the method are somewhat questionable, see below.\n\nThe method uses an augmented lagrangian objective to balance between reward and safety cost, but these methods are known for instability. Can the authors write more about how stable this process is,  and what they need to do to make it stable, etc.?", "methods_and_evaluation_criteria": "Twisted imagination: This procedure intends to generate trajectories with privileged information, by unrolling both privileged and unprivileged world models from the same starting state, using the same policy, and recording the pairs of state $(s^+_t, s^-_t)$. However, because each world model is stochastic, it seems unlikely that the trajectories are actually paired. Imagine an environment with a fork at the start state, and path 1 gives 100 reward while path 2 gives -100. The first WM predicts path 1, and the second WM predicts path 2.  So now your $(s^+_t, s^-_t)$ pairs should not really be paired together, since even though they are at the same timestep, they correspond to different locations in the environment. \n\nRepresentation alignment: Are you pulling  both state encoders towards each other? Or is it just one directional? \n\nSome additional baselines could be used, see below.", "theoretical_claims": "I read the theoretical section but did not check it carefully.", "experimental_designs_or_analyses": "I am not an expert in Safe RL, but there might be some essential Safe RL baselines that are missing, like Constrained Policy Optimization, TRPO-Lagrangian, etc.? In the main paper. Luckily, I skimmed the appendix and there seems to be comparison in section D.3. I think this should be moved to the main results.\n\nNext, for privileged reinforcement learning, a simple baseline is to train a privileged teacher with RL, and then distill it into an unprivileged student. I think this simple privileged baseline is important to show, to justify the additional complexities of using a model-based privileged RL approach.\n\nAll the environments seem to be done in low-dimensional state based environments. It is worth mentioning as a limitation, and image-based environments would be good future work.", "relation_to_broader_scientific_literature": "It can be seen as a model-based approach to leveraging privileged information in safe RL. \nPrevious work in privileged RL, to my knowledge, did not focus on safe RL. Although I would be sure that people have tried exploiting privileged information for safe RL, especially in robotics.", "essential_references_not_discussed": "The authors do a fair job of surveying privileged reinforcement learning with world models. \n\nThe algorithm seems quite similar to Scaffolder (ICLR24), which is mentioned and compared against. But could the authors give a more detailed comparison on the similarities and differences? It seems like PIGDreamer removes some parts and replaces some objectives.\n\nAnother privileged MBRL method not mentioned is Wasserstein Believer (ICLR24).", "other_strengths_and_weaknesses": "I think the application of privileged information towards solving Safe RL tasks is nicely motivated.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "Could you address my questions about the design choices? \n\nAblation experiments on the twisted imagination could be interesting, you could replace twisted imagination with the way Scaffolder does imagination?\n\nRunning a simple privileged RL baseline like distillation would be helpful for convincing readers to use your method over something simpler.\n\n> Compared to previous works (Lambrechts et al., 2024; Hu et al., 2024) that directly reconstruct\nprivileged information ...this method is significantly more robust when the privileged information is excessively informative for reconstruction\n\nI don't get what the authors are saying here, could you clarify?\n\nThere are no qualitative descriptions or videos of the tasks I can see. It would be nice to highlight strengths and weaknesses of each method's policies, by showing their rollouts and comparing them.\n\n---\nOverall: For now, I will lean on the positive side, assuming the authors take my feedback and address it.", "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": "hC5DZxnTJL", "reviewer": "Reviewer_ou54", "summary": "The work focuses on the development of a novel safe model-based reinforcement learning method. The researchers utilize the concept of privileged information and introduce the so-called asymmetric constrained partially observable Markov decision process (ACPOMDP) task, which requires access to the actual state in addition to observations for training the critic. They utilize the basic architecture of Dreamer and extend it to function simultaneously with both observations and actual states from the environment. Two models of the world are created: a naive model and a privileged model. During the deployment phase, only the naive model with access to observations is utilized. During the learning process, a multi-component loss is employed, incorporating privileged information. Safety is ensured through standard regularization using a Lagrangian approach. Experiments are conducted using Safety Gymnasium, and the proposed method demonstrates the same performance compared to SafeDreamer, Safe-SLAC, and other privileged approaches.\n\n## update after rebuttal\nThe authors have conducted additional experiments in a new Guard setting, enhancing their work's experimental part. These experiments show that PIGdreamer has low quality at a cost. These and previous results do not conclusively prove a significant advantage over other baselines for me. In my opinion, safe reinforcement learning methods should first demonstrate excellent results in terms of cost and only secondarily show good performance and high rewards. Considering my other comments, I am inclined to leave my assessment unchanged for now, but I will not object to the acceptance of the article.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the valuable suggestions provided by Reviewer OU54. In response, we have conducted additional experiments and clarifications, with the hope that these efforts will address your concerns.\n\n**Q1. Clarification Regarding Model Parameters and Deployment.**\n\nThe naive and privileged world models are trained jointly using the same set of parameters. During the deployment phase, the parameters for both models are loaded; however, only the naive world model and the actor are employed to generate actions.\n\n**Q2. Relation Between Our Theoretical Claim and Proposed Method.**\n\nWe appreciate your question. Our claim that a critic can develop more accurate value functions through direct access to the underlying states motivates our methodology in two key aspects:\n\n1. **Enhanced Value Estimation.** We utilize an asymmetric architecture to facilitate more precise value estimations, thereby improving the policy.\n2. **Mitigation of Risk Underestimation.** The analysis presented in Theorem 3.3 illustrates that reliance on partial observations results in the critic underestimating associated risks. This finding compels us to integrate privileged information into Safe Reinforcement Learning (SafeRL).\n\n**Q3. Guard Benchmark.**\n\nThank you for your suggestion. In response, we conducted additional experiments within the Guard benchmark under a limited timeframe. The Guard benchmark is a Safe Reinforcement Learning (SafeRL) framework that includes more complex agents, such as Ant and Humanoid, as well as more challenging tasks.\n\nIn our experiments, we utilized 64 x 64 pixel images as observations, along with low-dimensional sensors as privileged information. Our empirical results demonstrate that our method significantly outperforms other baseline approaches. Furthermore, we found that, in comparison to the Safety-Gymnasium benchmark, our method achieves highly competitive results within the Guard benchmark. We attribute this to the relative ease of the Safety-Gymnasium benchmark, which may obscure the advantages of our method. \n\n| Model |  | Goal_Ant_8Hazards |  |  | Goal_Ant_8Ghosts |  |  | Goal_Humanoid_8Hazards |  |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n|  | Training time | Cost | Reward | Training time | Cost | Reward | Training time | Cost | Reward |\n| SafeDreamer | 14.27 | 1.00 | 2.54 | 16.7 | 0.93 | 10.67 | 13.94 | 10.8 | 0.34 |\n| Informed-Dreamer(Lag) | 17.03 | **0.03** | 1.32 | 18.28 | 2.54 | 10.09 | 14.03 | 10.72 | -1.60 |\n| Scaffolder(Lag) | 24.82 | 0.04 | 1.86 | 27.14 | **0.83** | 6.15 | 23.8 | 13.14 | **2.13** |\n| Distill | 20.12 | 10.52 | 1.32 | 23.21 | 1.02 | 0.50 | 21.04 | 24.07 | 1.41 |\n| PIGDreamer(Ours) | 18.45 | 0.92 | **14.18** | 20.86 | 2.28 | **15.76** | 18.12 | **10.56** | 2.09 |\n\n**Q4. The performance improvement of our method in the Safety-Gymnasium benchmark.**\n\nThank you for your question. We acknowledge that our performance improvement in the Safety-Gymnasium benchmark may appear modest; however, there are several contributions of our work that we would like to highlight:\n\n1. **Privileged Learning for Safety.** We investigated the application of privileged MBRL in safety-critical tasks, a topic that is rarely addressed within the SafeRL community. In this regard, we reimplemented the most advanced privileged MBRL methods, Scaffolder and Informed-Dreamer, to incorporate safety constraints. The benefits of integrating privileged information are convincingly evidenced by our experimental results.\n2. **Efficiency.** As noted by other reviewers, our method achieves superior performance at the expense of a modest increase in training time, rendering it a practical solution for real-world applications.\n3. **More competitive results.** In **Q3**, we discuss the empirical results from the Guard benchmark, which we hope will address your concerns regarding the performance of our method.\n\n**Q5. Other Comments.**\n\nWe will update our manuscript based on your valuable suggestions. Additionally, we are in the process of developing a website for this paper, where we will include the code and videos.", "claims_and_evidence": "The statements proposed by the authors in the article on the effectiveness of using privileged information are generally confirmed by the experiments conducted.", "methods_and_evaluation_criteria": "The authors use the standard Safety Gymnasium benchmark, considering many tasks there, which is the standard in this field. In general, other benchmarks could be considered, such as SafeAntMaze based on MuJoCo.", "theoretical_claims": "The paper makes a fairly obvious claim that, in the presence of complete state information, a critic will form more accurate value functions than if it has incomplete observations. However, this claim does not add any significant insight to the approach. Instead, the proof in the appendix follows the logic presented in Pineau et al. (2006).", "experimental_designs_or_analyses": "The experiments were conducted according to the necessary requirements, but it is difficult to say that they were convincing. PIGDreamer does not differ significantly from the baseline with privileged information in terms of reward plots and cost value. However, the disadvantage of SafeDreamer in terms of not accessing true states is obvious. Figure 4 shows that all results are within the range of variations.", "relation_to_broader_scientific_literature": "In general, the proposed idea combines well-known methods: the SafeDreamer approach with a Lagrangian, and the idea of duplicating observations and states when training a world model. The positive aspect is that the authors have successfully formalized and trained all these elements into a single model. However, the experimental results do not seem very convincing.", "essential_references_not_discussed": "Most of the necessary papers are correctly cited.", "other_strengths_and_weaknesses": "The authors do not clearly explain how the naive and privileged world models  are separated by parameters. Based on the notation, it seems that they are trained with the same set of parameters, which is not well described. I have no other significant comments at this time.", "other_comments_or_suggestions": "The authors were careless about the links to pictures and pictures. The order is clearly messed up. The link to page 8 doesn't lead anywhere at all. Figure 4 shows strange x-axis signatures, especially on the cost graph.", "questions_for_authors": "Do naive and privileged world models have the same set of parameters? How is the privileged part discarded during the deployment phase?", "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": "mVzmiM1gm4", "reviewer": "Reviewer_Rrit", "summary": "This paper introduces PIGDreamer, a novel model-based reinforcement learning approach designed to enhance safety and performance in partially observable environments by leveraging privileged information during training. The authors propose Asymmetric Constrained Partially Observable Markov Decision Processes (ACPOMDPs), a theoretical framework that extends CPOMDPs by allowing the value function to access underlying states, thereby reducing the representation space and improving policy learning. PIGDreamer integrates privileged information through three key mechanisms: privileged representations, which align the state representations of the naive world model with privileged information; privileged predictors, which enhance the accuracy of reward and cost predictions; and privileged critics, which refine policy estimations. Empirical evaluations on the Safety-Gymnasium benchmark demonstrate that PIGDreamer significantly outperforms existing methods in terms of safety and task performance, achieving near-zero-cost performance while maintaining high rewards. The approach also exhibits superior training efficiency, with only a modest increase in training time compared to its unprivileged variant. Overall, PIGDreamer represents a significant advancement in model-based safe reinforcement learning for partially observable environments, effectively utilizing privileged information to achieve high performance and safety.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer Rrit for your thoughtful comments and questions. Your insights significantly contribute to the further enhancement of our paper.\n\n**Q1. A more comprehensive discussion of related work.**\n\nThank you for your suggestion. In the next version of our manuscript, we will discuss recent advancements in safe reinforcement learning that utilize privileged information.\n\n**Q2. Guard Benchmark.**\n\nWe appreciate your valuable suggestion. We conducted additional experiments within the Guard benchmark. Our empirical results indicate that our method significantly outperforms other baseline approaches. A more comprehensive response is available in **Q3** of our reply to Reviewer ou54.\n\n| Model |  | Goal_Ant_8Hazards |  |  | Goal_Ant_8Ghosts |  |  | Goal_Humanoid_8Hazards |  |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n|  | Training time | Cost | Reward | Training time | Cost | Reward | Training time | Cost | Reward |\n| SafeDreamer | 14.27 | 1.00 | 2.54 | 16.7 | 0.93 | 10.67 | 13.94 | 10.8 | 0.34 |\n| Informed-Dreamer(Lag) | 17.03 | **0.03** | 1.32 | 18.28 | 2.54 | 10.09 | 14.03 | 10.72 | -1.60 |\n| Scaffolder(Lag) | 24.82 | 0.04 | 1.86 | 27.14 | **0.83** | 6.15 | 23.8 | 13.14 | **2.13** |\n| Distill | 20.12 | 10.52 | 1.32 | 23.21 | 1.02 | 0.50 | 21.04 | 24.07 | 1.41 |\n| PIGDreamer(Ours) | 18.45 | 0.92 | **14.18** | 20.86 | 2.28 | **15.76** | 18.12 | **10.56** | 2.09 |\n\n**Q3. The assumptions of the ACPOMDP framework.**\n\nThank you for your question. In the ACPOMDPs framework, we assume that the privileged information consists of the underlying states of the environment; however, this assumption is impractical in real-world scenarios, as defining the underlying states and determining how to obtain them is challenging. Consequently, we relax this assumption to incorporate certain sensors that are available and precise, albeit costly, in real-world scenarios.\n\nFor example, in the Safety-Gymnasium benchmark, we define the privileged information as including the proprioceptive state, historical actions, and the relative positions of hazards and the goal. In the Guard benchmark, the privileged information consists of the proprioceptive state and LiDAR perceptions of hazards and the goal. All the aforementioned sensors are accessible in real-world applications, making our method a viable solution for practical implementation.\n\n**Q4. Detailed Statistical Analysis.**\n\nThank you for your suggestion. We apologize for our inability to complete all the experiments within the limited time available during the rebuttal period. Consequently, we intend to conduct a more comprehensive statistical analysis in the upcoming round of rebuttals.", "claims_and_evidence": "The claims made in the submission are generally well-supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Yes, the paper's proposed methods and evaluation criteria are well-suited for the problem and application.", "theoretical_claims": "I did not find any obvious issues in the proofs for theoretical claims in the paper.", "experimental_designs_or_analyses": "I did not find any obvious issues in the experimental designs or analyses in the paper.", "relation_to_broader_scientific_literature": "World Models: The use of world models to enhance sample efficiency and enable agents to learn from partial observations has been explored in several works (Hafner et al., 2019; Hogewind et al., 2022). These models learn environmental dynamics and task-specific predictions from past observations and actions.\n\nThe idea of using privileged information during training to improve performance has been explored in model-free RL (Hu et al., 2024; Lambrechts et al., 2024). However, its application in model-based RL, especially for safety-critical tasks, is less explored.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**Strengths**\n\nThe paper introduces the Asymmetric Constrained Partially Observable Markov Decision Processes (ACPOMDPs) framework, which is a novel extension of existing POMDPs. This framework allows the value function to access underlying states, leading to more efficient critic updates and superior policies. This theoretical contribution is significant and provides a new perspective on leveraging privileged information in partially observable environments.\n\nThe paper demonstrates strong empirical results on the Safety-Gymnasium benchmark, showing that PIGDreamer outperforms existing methods in terms of both safety and task performance. Additionally, the method achieves these improvements with only a modest increase in training time, making it a practical solution for real-world applications.\n\n**Weaknesses**\n\nThe paper could benefit from a more comprehensive discussion of related work, particularly recent advancements in safe RL that also leverage privileged information. This would provide a broader context for the contributions and highlight the unique aspects of PIGDreamer.\n\nThe paper evaluates PIGDreamer on the Safety-Gymnasium benchmark, but additional experiments on more diverse benchmarks or complex environments would strengthen the claims. This would provide further evidence of the method's robustness and generalizability.\n\nWhile the paper reports aggregate metrics and confidence intervals, a more detailed statistical analysis, including significance tests, would further strengthen the empirical results. This would provide additional confidence in the robustness of the findings.\nThese points highlight the key strengths and areas for improvement in the paper, providing a balanced view of its contributions and potential enhancements.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Could you elaborate on the assumptions made in the ACPOMDPs framework, particularly regarding the nature of the privileged information and its impact on the generalizability of the results? How do these assumptions affect the applicability of your framework to real-world scenarios where privileged information might not always be available or reliable?\n\nIn addition to the Safety-Gymnasium benchmark, have you considered evaluating PIGDreamer on other benchmarks or real-world datasets? If so, could you share any preliminary results or insights from those evaluations?", "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:

Meta-review: The paper presents PIGDreamer, a model-based reinforcement learning approach that leverages privileged information during training to enhance safety and performance in partially observable environments. The authors introduce the Asymmetric Constrained Partially Observable Markov Decision Process (ACPOMDP) framework, which allows the value function to access underlying states, improving policy learning. PIGDreamer incorporates privileged information through privileged representation alignment, asymmetric actor-critic structures, and enhanced reward prediction mechanisms. The empirical results, tested on the Safety-Gymnasium benchmark, indicate that the approach achieves superior safety and task-centric performance compared to existing methods.

The paper proposes a novel approach to incorporate privileged information for safe RL, provides a theoretical foundation for privileged RL, and has promising experimental results in terms of safety, task performance, and efficiency compared to baselines. 

The gains are small, the experimental benchmark tasks are somewhat non-standard, and the qualitative presentation of results is poor. Some of these were addressed in the rebuttals, leading to overall positive scores from each reviewer. The authors are encouraged to incorporate these changes and new results into future versions.
Decision: Accept (poster)