PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning

We thank Reviewer Rrit for your thoughtful comments and questions. Your insights significantly contribute to the further enhancement of our paper.

Q1. A more comprehensive discussion of related work.

Thank you for your suggestion. In the next version of our manuscript, we will discuss recent advancements in safe reinforcement learning that utilize privileged information.

Q2. Guard Benchmark.

We 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.

Q3. The assumptions of the ACPOMDP framework.

Thank 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.

For 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.

Q4. Detailed Statistical Analysis.

Thank 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.

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.

Q1. Clarification Regarding Model Parameters and Deployment.

The 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.

Q2. Relation Between Our Theoretical Claim and Proposed Method.

We 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:

1. Enhanced Value Estimation. We utilize an asymmetric architecture to facilitate more precise value estimations, thereby improving the policy.
2. 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).

Q3. Guard Benchmark.

Thank 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.

In 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.

Q4. The performance improvement of our method in the Safety-Gymnasium benchmark.

Thank 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:

1. 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.
2. 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.
3. 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.

Q5. Other Comments.

We 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.

We would like to express our gratitude to Reviewer byT1 for your insightful comments.

Q1. The mismatch arises from the Twisted Imagination (TI) and the ablation of TI with the Nested Latent Imagination (NLI).

Thank 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^{+}$ .

Q2. The implementation of Representation Alignment.

In our work, we pull both state encoders towards each other using the method proposed in Dreamerv3.

Q3. Clarification regarding the privileged representation.

In 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.

Q4. Comparison with Scaffolder.

Thank 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:

1. Theoretical Foundation: We propose our method within the ACPOMDP framework while Scaffolder has made limited advancements in this area.
2. Privileged Representations: see Q3.
3. 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.
4. Privileged Imagination: see Q1.
5. Privileged Exploration: We eliminate this component for the following reasons:
   1. It prolongs the training time.
   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.
6. Privileged Critics: same with Scaffolder.
7. Efficiency: Our method demonstrates enhanced efficiency in utilizing privileged information.

Q5. The teacher-student distillation baseline.

Thank 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.

Q6. Experimental setup.

Thank 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.

Q7. Wasserstein Believer.

We plan to incorporate a discussion of its relevance in the next version of our paper.

Q8. The stability of the augmented lagrangian method.

We 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.