Uncertainty-Sensitive Privileged Learning

W1

Theoretical justification.

We are grateful to the reviewer for this insightful comment regarding the theoretical justification for our proposed framework. Heeding this suggestion, we have developed a formal analysis to address this concern. Specifically, by combining a Lipschitz assumption on the policy with a result from Lemma 4 in [1], we derive the following cumulative error bound that formally connects prediction uncertainty to the behavioral discrepancy between the PP and DP.

Let the policy in USPL be denoted as $\pi(\cdot\mid o_ c, o_ p, \sigma_ p)$, where $o_ c, o_ p,$ and $\sigma_ p$ represent the common observation, privileged observation, and uncertainty estimation, respectively. The key difference between the PP and the DP lies in the source of $o_ p$. Let the privileged observation input to the PP at timestep $h$ be $o_ p^h$ (the ground truth) and the one input to the DP be $\hat{o}_ p^h$ (the prediction). We assume there exists an error bound $\epsilon^h$ such that:

$$\Vert o_ p^h - \hat{o}_ p^h \Vert_ 2 \leq \epsilon^h.$$

Here, $\epsilon^h$ represents the upper bound on the prediction error of the privileged observation. We assume that the policy $\pi(\cdot\mid o_ c, o_ p, \sigma_ p)$ is Lipschitz continuous with respect to the privileged observation $o_p$:

$$D_ {TV}(\pi(\cdot\mid o_ c, o_ p^h, \sigma_ p), \pi(\cdot\mid o_ c, \hat{o}_ p^h, \sigma_ p)) \leq L(o_ c, \sigma_ p) \Vert o_ p^h - \hat{o}_ p^h \Vert_ 2,$$

where $L(o_ c, \sigma_ p)$ is the Lipschitz constant of $\pi$ with respect to $o_ p$ for a given $o_ c$ and $\sigma_ p$. This constant will be large if the policy is highly sensitive to privileged observations. We assume $\epsilon^h$ is known and is used as the uncertainty input to the policy. Let $\pi_ h^{PP} = \pi(\cdot\mid o_c, o_ p^h, \epsilon^h)$ and $\pi_ h^{DP} = \pi(\cdot\mid o_ c, \hat{o}_ p^h, \epsilon^h)$. Under these assumptions, the following inequality holds for the cumulative discrepancy between the state visitation distributions:

$$\frac{1}{H}\sum_ {h=1}^H D_ {TV}(d_ h^{\pi^{PP}}, d_ h^{\pi^{DP}}) \leq \sum_ {h=1}^H \mathbb{E}_ {o_ c \sim d_ h^{\pi^{PP}}}\left[L(o_ c, \epsilon^h)\epsilon^h\right],$$

where $d_ h^{\pi^{PP}}$ and $d_ h^{\pi^{DP}}$ are the visitation distributions of $o_ c$ at step $h$ for the PP and DP, respectively, and $H$ is the horizon. A detailed derivation of this bound will be provided in the revised paper.

Relationship between the Error Bound and our USPL Framework:

This bound formally demonstrates that the cumulative behavioral discrepancy is upper-bounded by a function of the privileged observation prediction error ($\epsilon^h$) and the policy's sensitivity to that information ($L(o_c, \epsilon^h)$). To minimize this compounding error, one can pursue two primary strategies:

• (S1) Minimize the prediction error $\epsilon^h$: This involves improving the prediction accuracy of the observation encoder.
• (S2) Reduce the policy's sensitivity to the privileged observation: This involves decreasing the Lipschitz constant $L(o_c, \epsilon^h)$.

In USPL, S1 is directly addressed by the optimization objective of our observation encoder and reward transformation, both of which aim to reduce $\epsilon^h$. S2 is realized through privileged blurring, which adds noise to the privileged observation during training to reduce the policy's over-reliance on it. However, at the beginning of an episode, when the agent's perception of the environment is insufficient, $\epsilon^h$ is typically large. At the same time, solving the task often requires leveraging the privileged observation, which means $L(o_ c, \epsilon^h)$ also needs to be large. To prevent the compounding error from escalating due to both $\epsilon^h$ and $L(o_ c, \epsilon^h)$ being high simultaneously, USPL also achieves an adaptive dependence on privileged observations by taking uncertainty as policy input:

• When uncertainty ($\approx \epsilon^h$) is high, the privileged observation is very noisy. Thus, the policy learns to distrust the privileged input (effectively reducing $L(o_ c, \epsilon^h)$). Guided by the transformed reward, it prioritizes exploration to gather information, thereby reducing future uncertainty.
• As uncertainty ($\approx \epsilon^h$) decreases, the privileged observation becomes more reliable. The reward for task completion is magnified by reward transformation. As a result, the policy learns to increase its reliance on the privileged observation (increasing $L(o_c, \epsilon^h)$) to efficiently solve the task.

This coordinated control of both $\epsilon^h$ and $L(o_c, \epsilon^h)$ effectively tightens the upper bound on the cumulative error. This, in turn, minimizes the behavioral discrepancy between the PP and DP, mitigating performance degradation at deployment.

Empirical Evidence: As demonstrated by our experiments with manually set uncertainty (Figure 7, Figs. 14-17 in the appendix, and videos in the supplementary material), USPL policy successfully learns this adaptive capability. When provided with high uncertainty, the policy indeed refrains from exploiting the privileged observation—even when the prediction is accurate—and instead prioritizes exploration. It only begins to leverage the privileged observation to complete the task once the uncertainty is reduced, validating our theoretical analysis.

We will add a new subsection in Sec. 3 of the revised paper to elaborate on this theoretical analysis, thereby strengthening the paper's rationale.

[1] Understanding adversarial imitation learning in small sample regime: A stage-coupled analysis. arXiv preprint arXiv:2208.01899.

Q1

The choice of a uniform distribution to randomize the privilieged observations.

We apologize for the confusion. The form of noise perturbation in Eq. (3) is indeed a hyperparameter. We found that when using Gaussian noise ($\mathcal{N}(\cdot\mid 0, \sigma_p^2)$), even with a high fixed uncertainty $\sigma_p$, the PP could sometimes still complete the task without performing information-gathering actions. This led to a significant behavioral discrepancy between the PP and the DP. On the other hand, increasing the noise magnitude using a distribution like $\mathcal{N}(\cdot\mid 0, (K\sigma_p)^2)$ with $K>1$ often resulted in perturbed observations $\hat{o}_p$ that fell outside their practical range and the output range of the observation encoder's mean $\mu_p$. These are values that would never be encountered during deployment, and their frequent occurrence could potentially reduce training efficiency. Moreover, in our experiments, the environment parameters corresponding to $o_p$ are sampled from a bounded uniform distribution. By using uniform noise for perturbation, we can adjust the distribution's bounds to ensure that the final $\hat{o}_p$ remains within a valid range. The $\hat{o}_p$ with the maximum perturbation follows the prior distribution of $o_p$. We will incorporate the aforementioned discussion into Sec. 3.2 of the revised version of the paper.

The table below compares the success rate using the perturbation method from Eq. (3) versus a Gaussian distribution $\mathcal{N}(\cdot\mid 0, (K\sigma_p)^2)$.

In the Signpost Nav task, the performance difference is minor. However, in the Biased Quadrotor task, the performance with Gaussian noise is substantially worse.

Q2

The anomalous success rate standard error for the Biased Quadrotor in Table 1.

Thank you for pointing this issue out. This was indeed a calculation error on our part. The success rates for the 5 seeds were [0.983, 0.955, 0.956, 0.963, 0.977], with a mean of approximately 0.967 and a standard error of 0.0056. We mistakenly forgot to multiply the standard error by 100, which led to the anomalous value in the table. The correct success rate is $96.7 \pm 0.6$. We will correct this in the revised version of our paper.

---

We are truly grateful for your thoughtful and constructive feedback. Your suggestions were incredibly helpful and have guided us in making significant improvements to the paper. We believe the paper is much stronger now as a result. We hope our responses have thoroughly addressed all your concerns, and we look forward to your feedback. Please do not hesitate to contact us if any points require further discussion.

W1

Could the authors more clearly explain how does their work differ from [1] and why it is novel?

[1] Meta reinforcement learning as task inference, Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A. Ortega, Yee Whye Teh, Nicolas Heess, 2019

We thank the reviewer for pointing out this relevant work. USPL is fundamentally distinct from [1] as USPL is a privileged learning (PL) framework, while [1] is not. This core difference leads to distinct training methodologies, primary challenges, and novel contributions. We highlight three key distinctions:

1. Usage of privileged observations

• USPL: During training, the Privileged Policy (PP) has direct, real-time access to privileged observations. This allows the PP to focus on completing the task without needing to learn active information-gathering behaviors, leading to a more efficient RL training process. Our primary focus is then on bridging the "behavioral divergence" between this PP and a deployment policy that operates on target observations.
• Method in [1]: The policy never directly uses privileged information as an input. Instead, privileged data is used solely to supervise an auxiliary loss for a belief encoder. Consequently, their single policy must learn complex, active information-gathering behaviors to infer the missing information, making the policy training significantly more challenging.

2. Core technical contribution

• USPL: A central contribution of our work is the prediction and utilization of uncertainty. We explicitly model the uncertainty of the privileged observation prediction and leverage this uncertainty to mitigate behavioral divergence.
• Method in [1]: The concept of prediction uncertainty is entirely absent. Their work focuses on task inference through belief state representation without considering the reliability of that belief.

3. Empirical comparison

• Our baseline "PrivRecons" in Table 1 is conceptually analogous to the approach in [1]. PrivRecons uses a observation encoder trained via supervised learning to predict privileged observations, and a single policy uses this encoder's output for both training and deployment.

• The experimental results clearly show that USPL substantially outperforms PrivRecons. This result empirically validates the benefits of our privileged learning framework and our novel use of uncertainty, highlighting the advantages over the methodology presented in [1].

We hope this comparison clarifies the distinctions of our approach and its contributions relative to [1]. We will include the comparisons with [1] in the revised version of our paper.

W2

The experiments are still limited to low-dimensional privileged observations.

This is indeed a known limitation of the current work. Extending USPL to scenarios with high-dimensional, redundant, and time-varying privileged observations is indeed a challenging and important direction for future research. Without additional techniques, the performance of USPL in such settings would heavily depend on the capacity and learning efficiency of the observation encoder.

A promising approach to address this would be to integrate dimensionality reduction techniques. For instance, [2] and [3] introduced a privileged encoder that maps high-dimensional privileged observations into a low-dimensional latent space, and they have demonstrated excellent results. Similarly, USPL could adopt such a method to first project the privileged observations into a compact latent space and then train the observation encoder to predict this lower-dimensional representation.

We will continue to conduct more in-depth research on this limitation in our future work. We will also extend the discussion of this limitation in the revised version of our paper.

[2] RMA: Rapid motor adaptation for legged robots. RSS 2021.

[3] Deep whole-body control: Learning a unified policy for manipulation and locomotion. CoRL 2022.

---

We would like to express our sincere gratitude once again for your insightful and constructive suggestions. We hope that our response and the proposed revisions have thoroughly addressed your concerns. We believe the paper is significantly stronger as a result of your guidance, and we would be grateful if you would consider these clarifications in your final assessment. Please do not hesitate to let us know if any further questions arise.

Dear Reviewer mnJ9,

We would like to extend our sincerest gratitude for your incredibly insightful and thoughtful review. Your constructive comments have been invaluable, providing us with a clear path to significantly strengthen our manuscript.

We hope that our rebuttal and planned revisions have thoroughly addressed your concerns. As the discussion period is drawing to a close, we wanted to reach out one more time. We remain fully available and would be genuinely delighted to provide any further clarification you might need.

If our responses and revisions have successfully resolved your initial concerns, we would be deeply grateful if you would consider these improvements in your final evaluation of our work.

Thank you once again for your invaluable guidance and engagement in this process.

Respectfully,

The Authors

W1

Lack of comparison in terms of the computation cost with the baseline methods.

Thank you for your valuable feedback. We agree that a discussion on computational cost is crucial for a complete comparison. In our experiment, all baseline methods were implemented with identical network architectures to ensure a fair comparison. The differences between the baselines lie in the loss functions used for training various modules. We observed that the primary computational overhead stems from environment interaction and network updates, rather than the specific loss calculations. Consequently, the overall computational costs for different methods are quite similar.

We show the average training time per epoch for each method on the Signpost Nav task in the following:

As shown, RESeL-PPO, which trains the entire DP network end-to-end with a PPO loss, is approximately 13% faster than the other methods that require additional loss computations. The training times for USPL and the other three baselines show no significant differences. We will incorporate this comparison into the revised version of our paper.

W2

In the ablation studies, the performance improvement seems to largely come from other compoenents like blurring and modifying the reward, which make it questionable that what the main contribution of the propose method should provide. Is it possible to apply these components to the baseline methods and also significantly improves the performance?

We sincerely thank the reviewer for this insightful question. You have pinpointed a crucial aspect of our ablation study that merits a more detailed explanation:

1. We first want to address the question of whether privileged blurring and reward modification could be applied to baseline methods. These components are not standalone techniques, and are fundamentally reliant on the uncertainty estimation $\sigma_p$, which is the central contribution of USPL. They cannot be applied to baseline methods that do not estimate this uncertainty. We also recognize, with thanks to your comment, that our original table caption was a source of confusion. The "w/o uncert" variant does, in fact, still compute and use uncertainty for blurring and reward modification. It only omits uncertainty as a direct input to the PP. We will rename this variant to prevent future ambiguity.

2. Regarding the performance gain comparison of different components, we would like to clarify that our ablation results should be viewed not as a competition between components, but as a demonstration of how they play distinct and complementary roles in a multi-stage process.

   An optimal DP would ideally adapt its behavior online based on its current Progress of Collecting Privileged Information (PCPI), which is inversely correlated with our uncertainty estimation $\sigma_p$. Our framework aims to build such a policy based on PP, and this requires satisfying two essential conditions:

   • First, the PP must explore and learn the optimal behaviors for different levels of PCPI. During training, our privileged blurring and reward transformation components work in concert to achieve this foundational goal. Privileged blurring simulates states of high uncertainty, forcing the policy to learn how to act when privileged information is unreliable. Concurrently, reward transformation incentivizes the policy to actively explore and discover actions that reduce uncertainty (i.e., increase PCPI). The necessity of this exploration stage is evident in Table 2: removing either privileged blurring or reward transformation leads to a collapse in the success rate, proving they are essential for learning the necessary behaviors.

   • Second, the PP must be able to dynamically select the appropriate behavior online based on its current PCPI. This is accomplished by providing the uncertainty $\sigma_p$ as a direct input to the policy. The uncertainty serves as a real-time signal, allowing the policy to intelligently switch between information gathering (when uncertainty is high) and task execution (when uncertainty is low). However, this online adaptation is only effective if the optimal behaviors have already been discovered via privileged blurring and reward transformation. This explains the poor performance of only uncert in Table 2.

   Incorporating uncertainty as a policy input is crucial for achieving optimality and reliability. To quantify its importance, we analyzed the task failure rate, a key metric for evaluating policy reliability. The results are presented in the table below. As the table shows, even with the powerful exploration capabilities provided by blurring and reward transformation, the policy still has a notable average failure rate of 9.4%. Providing uncertainty as a direct policy input reduces this failure rate by 80.9%, bringing it down to just 1.8%.

   This result powerfully illustrates that while blurring and reward shaping are essential for teaching the policy what to do, the uncertainty input is the crucial final piece that tells the policy when to do it. It is this ability to dynamically adapt based on a direct measure of its own uncertainty that makes the policy reliable, and ultimately more successful.

   Task	Failure Rate (w/o Uncert. Input)	Failure Rate (w/ Uncert. Input)	Failure Rate Reduction
   Signpost Nav	3.0%	2.2%	26.7%
   Square Maze	5.3%	1.5%	71.7%
   Lateral Choice	8.0%	1.2%	85.0%
   Midpoint Choice	21.1%	2.3%	89.1%
   Average	9.4%	1.8%	80.9%

   ---

   We are truly grateful for your perceptive feedback. It prompted a deeper analysis that has, we believe, significantly improved the paper. We hope our explanation clarifies the unique role of each component within our framework. We sincerely look forward to any further discussion and hope that these clarifications and new results will lead you to view our contributions even more favorably.

Dear Reviewer GVJr,

Thank you for your thoughtful and constructive review. Your feedback was incredibly helpful and has allowed us to significantly improve our manuscript.

We have incorporated your suggestions and hope our revisions and rebuttal have fully addressed your concerns. We welcome any further discussion and would be happy to clarify any points as the deadline approaches. We hope you find the paper is much improved and would be grateful if you would consider these changes in your updated assessment.

Thank you again for your time and valuable guidance.

Respectfully,

The Authors

W1

Lack of sensitivity analysis to hyperparameters.

Thank you for raising this important point. Our main hyperparameters include those that affect the reward transformation ($C_{\min}$, $C_{\max}$) and the hyperparameter $K$, which controls the magnitude of the perturbation noise to $o_p$. We present the results of our sensitivity analysis for these two types of hyperparameters.

First, we conduct an sensitivity study on $C_{\min}$ and $C_{\max}$, which affect the scale coefficient for the original reward. If $\log(\sigma_p) > C_{\min}$, the scale coefficient reaches its minimum. Conversely, the scale coefficient increases as $\sigma_p$ decreases, until $C_{\min} - \log(\sigma_p) > C_{\max}$, at which point the coefficient reaches its maximum and stops increasing. Therefore, $C_{\min}$ and $C_{\max}$ jointly determine the range of the scale coefficient. We tested different combinations of $C_{\min}$ and $C_{\max}$ in both the Signpost Nav and Biased Quadrotor tasks, with the results shown in the table below.

For the Signpost Nav task, the first three rows show that a relatively small range for the scale coefficient limits the impact of the reward transformation, resulting in a task success rate below 90%. As the range expands, the success rate consistently stays above 97%. The Biased Quadrotor task exhibits a similar trend at the lower end of the range, with the narrowest range resulting in a very low success rate ($12.5\%$). However, it also reveals that an overly large range can be detrimental. As shown in the last two rows, when the range becomes too wide, the performance drops significantly to $14.0\%$ and $53.1\%$, respectively.

Overall, the results suggest that as long as the range of the scale coefficient is not excessively large or small, the policy's success rate is largely insensitive to the specific values of $C_{\min}$ and $C_{\max}$. Within an appropriate range, different combinations of these parameters are capable of achieving near-optimal performance.

Next, we perform a sensitivity analysis on $K$, which affects the magnitude of the perturbation noise to $o_p$. The results are in the table below. The results in Signpost Nav show that the value of $K$ had almost no impact on the policy's success rate. As for the Biased Quadrotor task, the results show that a smaller value of $K$ leads to a significant drop in the success rate, whereas for $K \geq 10$, the policy consistently achieves a high success rate. USPL performs well across a wide range of values for $K$. After $K \geq 10$, the policy's success rate is largely insensitive to the value of $K$.

We will include these results and the corresponding discussion in the revised version of our paper.

W2

It would be nice to see some experiments on other applications of privileged learning.

Thank you for this constructive suggestion. We agree that demonstrating the applicability of our method beyond robotics environments would significantly strengthen the paper. As a result, we test our method on the challenging credit assignment task discussed in [1], Active T-Maze (with the credit assignment length set to 500). Active T-Maze is a grid-world task where the agent's navigable area is shaped like a "T" rotated 90 degrees clockwise (as in Fig. 1 of [1]). The goal is located at either the top-right or bottom-right corner. The optimal policy for the agent is to first move one step left to an "oracle" position to reveal the goal's location, then move right to a junction, and finally decide whether to move up or down based on the privileged information it received. [1] found that traditional Transformer-based and RNN-based RL methods struggle to solve such credit assignment tasks efficiently. We approach the Active T-Maze task from a privileged learning perspective, with the results shown below. In this experiment, USPL was trained with 5 different random seeds, while other baselines were trained with 3 random seeds due to limited computational resources.

As the results show, USPL achieves a perfect success rate on this highly challenging task, while the other baselines only achieve around a 50% success rate. We hypothesize that USPL's explicit uncertainty input and reward transformation significantly alleviate the difficulty of credit assignment. These results demonstrate the effectiveness of USPL on credit assignment problems. We will add these findings and discussion to the revised version of our paper.

[1] When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment. NeurIPS 2023.

Q1

Even if you opt for a proxy encoder, does it not need to be trained alongside PP?

You are right, the proxy encoder does need to be trained alongside the PP. However, this proxy encoder would be more lightweight than the observation encoder we currently use, leading to faster training. This is especially true for image-based inputs. Currently, training the PP requires processing images with the observation encoder, which is slow. By contrast, using a proxy encoder to process the generally lower-dimensional historical privileged observations would be much more efficient.

Q2

There is much room for improvement in the writing part. Sections 1 and 2 often sound repetitive.

Thank you for this valuable feedback. We will thoroughly revise the writing of the entire paper in the next version, with a particular focus on Sections 1 and 2, to improve clarity and readability.

---

Thank you once again for your detailed review and constructive suggestions. Your feedback has been crucial in helping us enhance the quality of our paper. In light of these clarifications and additional experiments, we would be grateful if you would consider a higher rating. We would be more than happy to provide any further clarifications if needed.

Dear Reviewer ru4E,

We would like to express our profound appreciation for your perceptive and thorough review. Your suggestions were truly instrumental, offering critical insights that have allowed us to substantially enhance the quality and clarity of our manuscript.

We hope that our rebuttal and the accompanying revisions have fully addressed your concerns. Mindful that the discussion window is concluding soon, we wanted to make ourselves available for any final questions you might have. We would be more than happy to engage further on any point. It is our sincere hope that these revisions, which were directly motivated by your feedback, have strengthened the paper to your satisfaction. We would be very grateful if you would take these developments into account in your updated assessment.

Thank you once more for your expert guidance and for your constructive partnership in refining this work.

Respectfully,

The Authors