Paper: Privileged Sensing Scaffolds Reinforcement Learning (ICLR-2024)

Reviewer comment (Writing):
Increasing the clarity around the Posterior and detailing how it is used to transition from the privileged latent state to the non-privileged latent state would greatly enhance understanding of the method.

Author response:
We noticed several reviewers also had questions about the methodology, in particular the Nested Latent Imagination where we unroll the scaffolded dynamics model with the target policy and posterior. The reviewer may be interested in our response to Reviewer xL3j on the motivation behind learning separate scaffolded and target representations. 

We have added additional details in Appendix A.3 explaining the posterior and its usage in Nested Latent Imagination (NLI) to the appendix, along with pseudocode and a side-by-side figure comparing Nested Latent Imagination with the standard Latent Imagination procedure from DreamerV3 which only uses the impoverished target world model to generate trajectories.

Reviewer comment (Experiments):
In the experimental design, the wrist camera and touch features don't appear to be excessively privileged or substantially different from the target observations. It would be beneficial to experiment with more oracle-like elements in the simulator as privileged sensory inputs. For instance, the oracle contact buffer or geodesic distance could be considered.

Author response:
Scaffolder can be used for sim-to-real, with simulator states acting as the privileged sensors. In this scenario, it makes sense to use oracle-like elements in the simulator. Some tasks in our S3 suite do indeed resemble such a setting: for example, the dexterous manipulation tasks use ground truth object poses as privileged sensors.

However, it’s important to note that Scaffolder generalizes privileged sensors from simulator state, to noisy, high-dimensional observations - enabling privileged training in the real world. Some of the S3 tasks are designed to emulate a real world learning setup and as a result use realistic privileged sensors such as wrist cameras or touch sensors.

Reviewer comment (Writing):
More clarification on the posterior and embedding component in the method part.

Author response:
We noticed several reviewers also had questions about the methodology, in particular the Nested Latent Imagination where we unroll the scaffolded dynamics model with the target policy and posterior. The reviewer may be interested in our response to Reviewer xL3j on the motivation behind learning separate scaffolded and target representations. 

We have added additional details in Appendix A.3 explaining the posterior and its usage in Nested Latent Imagination (NLI) to the appendix, along with pseudocode and a side-by-side figure comparing Nested Latent Imagination with the standard Latent Imagination procedure from DreamerV3 which only uses the impoverished target world model to generate trajectories.

Reviewer comment (Evaluation):
For some claims made in the paper, it’s actually not quite convincing. For “In other words, much of the gap between the observations o− and o+ might lie not in whether they support the same behaviors, but in whether they support learning them.”, some additional visualization like trajectory visualization might be helpful to strengthen the claim, since the similar reward score does not necessarily result in similar behavior.

Author response:
Thank you, we recognize now that this statement was not clearly worded. Target and privileged policies often behave very differently in our settings, of necessity: for example, consider Blind Pick, where the target policy must first grasp around the table to find an object (see [videos on the website](https://sites.google.com/view/sensory-scaffolding#h.qlmom6baw90)), whereas the privileged policy sees the object and can directly pick it up. We do not of course claim that such wildly different behaviors are the same, just that they might achieve similar task metrics (e.g. task success rates). We will reword to: “In other words, much of the gap between the observations o− and o+ might lie not in whether they support achieving similar task performance metrics, but in whether they support _learning_ them."

Reviewer comment (Evaluation):
For runtime comparison, since the speed of given GPUs varies a lot, it might be better to compare the wall-time with similar system configuration, assuming the wall-time is consistent across different seeds.

Author response:
We will update the appendix with wall times recorded on a similar system configuration.

Reviewer comment (Reproducibility):
Regarding some technical details, in section C1, it says “We launch 4-10 seeds for each method”, what’s the exact meaning of using different number seeds across methods or across environments?

Author response:
Due to computational limitations, it was expensive to run 10 seeds for all methods across all environments. We respectfully ask the reviewer to consider our large evaluation setup - 10 tasks and 5 model-based methods each taking a GPU and up to a day to run would result in 500 GPU days worth of compute.  

We have recorded the current number of seeds for each task below.
|  | Blind Pick | Blind Locomotion | Blind Deaf Piano | Blind Numb Pen | Blind Numb Cube | Noisy Monkey | Wrist Pick-Place | Occluded Pick-Place | RGB Pen | RGB Cube |
|---|---|---|---|---|---|---|---|---|---|---|
| Scaffolder | 8 | 5 | 5 | 5 | 5 | 5 | 5 | 5 | 5 | 5 |
| DreamerV3 | 10 | 10 | 5 | 5 | 5 | 5 | 5 | 5 | 4 | 4 |
| Informed Dreamer | 8 | 5 | 5 | 5 | 5 | 5 | 5 | 5 | 5 | 5 |
| DreamerV3 + BC | 5 | 5 | 5 | 5 | 5 | 5 | 5 | 5 | 4 | 4 |
| Guided Observability | 5 | 5 | 5 | 5 | 5 | 5 | 5 | 5 | 5 | 5 |
| RMA+ (100M) | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 |
| AAC (100M) | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 |

Note that our baselines RMA+ and AAC are run for 100M environment steps (20-400x the normal step budget), at which point we expect them to have overcome variations across initialization seeds. As a result, we just reported 1 seed for RMA+ and AMC.

We are running more seeds, and will update the results with 10 seeds for all tasks when ready, although this will likely happen after the rebuttal period. We do not anticipate any major changes in trends.

Reviewer comment (Experiments):
I agree that it makes sense to evaluate the proposed method on the newly proposed benchmark, for motivations mentioned in the paper. However, the paper would still benefit from evaluating extra existing benchmarks, just for reference.

Author response:
> Extra results on existing benchmarks just for reference?

Thank you for this suggestion. The majority of relevant prior works have been evaluated in a restricted setting where privileged sensors are always low-dimensional simulator states (e.g. AAC, RMA). In other cases, tasks were often much simpler: for example, consider the “Bumps” robotic task in [2], which is akin to a simpler version of the “Blind Pick” task in our S3 suite. In Bumps, a blind robot must feel around a tabletop to find randomly placed bumps using proprioception and touch, and the privileged observations reveal bump locations. We deliberately constructed S3 to have complex tasks such as dexterous manipulation, and high-dimensional / noisy privileged sensors like cameras rather than merely low-dimensional simulator states. As a result, S3 is, to our knowledge, the most general and most challenging suite of environments testing learning with privileged sensory information. 

However, we agree that a comparison on previously proposed tasks is valuable: we are setting up experiments in the Bumps environment described above, and will report on our progress before the rebuttal period closes. We expect broadly similar trends to Blind Pick since, as described above, these tasks have similar privileged and target sensor gaps.

[2] Nguyen, Hai, et al. "Leveraging fully observable policies for learning under partial observability." Conference on Robot Learning. PMLR, 2022.

Reviewer comment (Theory):
I am curious about one particular design choice. Why do the authors choose to predict target latent from privileged latent, bottlenecked by target observation. Why not use the same latent, shared by both the privileged and target modules?

Author response:
> Design choice: why predict target latent from privileged latent, bottlenecked by target observation,  why not use the same latent, shared by both the privileged and target modules?

Thank you for this question. We notice that several reviewers have had questions about this part of our approach, and will aim to expand and improve the exposition of Nested Latent Imagination, following the description below.

Scaffolder tries to simultaneously achieve the following two things: (1) it trains a scaffolded world model operating on scaffolded observations $o^+$ (represented into a “scaffolded latent” $z^+$) to more faithfully model the environment and generate better synthetic training experiences, (2) it trains a target policy, which in our setting, is restricted to operate only on information from target observations $o^-$, represented in a target latent $z^-$. In other words, $z^-$ cannot contain information that cannot be inferred from target observations alone. 

A shared representation $z^+=z^-$, as used in some prior approaches like RMA and Informed Dreamer, would violate these requirements: if it contained information from privileged sensors to enable better world model learning as required in (1), it would not be fully inferable from $o^-$ alone, as required by (2). Consider, for example, S3 tasks like Blind Picking or Blind Locomotion, which have target observations (e.g. proprioception) that are not predictive of privileged observations (images of objects). 

Scaffolder proposes a relatively straightforward solution to achieve both (1) and (2). It keeps $z^+$ and $z^-$ separate, but learns a target embedding translator $z^+ \rightarrow e^-$. This allows a mapping from $z^+ \rightarrow e^- \rightarrow z^-$ via the target posterior, so that the synthetic experiences generated in the scaffolded WM can still be used to train target policies. Note that the scaffolded observations $o^+=[o^-, o^p]$ include and extend the target observations, so it is reasonable to learn a unique mapping from $z^+$ to $e^-$.