Privileged Sensing Scaffolds Reinforcement Learning

Thank you for your review!

Scaffolder requires doing 2x the world model fitting, as both the z^- and z^+ models need to be fit for the approach to work. In general, this is "fair" for model-based RL, which is usually judged in number of environment interactions rather than number of gradient steps, but it very definitely is a more complex system and this can introduce instability. A common actor-critic criticism is that the rate of learning between the actor and critic needs to be carefully controlled such that neither overfits too much to the other. Scaffolders seems to take this and add another dimension for the rate of learning between the privileged world model and unprivileged one, as well as the learning speed of the privileged exploratory actor

We found the hyperparameter tuning process for Scaffolder, and DreamerV3 in general, to be short and easy. On a completely new environment and task, we found that DreamerV3 only needs tuning for two hyperparameters, the model size and the update to data (UTD) ratio. We follow an easy guideline for tuning these - more complicated dynamics generally require larger models, and tasks with harder exploration require more data and fewer updates (low UTD).

We did not tune these two DreamerV3 hyperparameters towards Scaffolder - rather, when we used model sizes and UTD ratios from when we ran DreamerV3 with privileged inputs as a reference method used for computing the upper bound scores. These same settings are then used for all DreamerV3 methods (Scaffolder, Informed Dreamer, DreamerV3+BC, Guided Observability) to be consistent.

As a result, we found hyperparameter settings that work for tasks with similar properties:

• Model / UTD
• Small / 512: Blind Pick, Blind Locomotion, Wrist Pick-Place, Occluded Pick-Place (Simple Dynamics, Easy Exploration)
• Large / 512: Blind Deaf Piano (Complex dynamics, Easy Exploration)
• Large / 16: Noisy Monkey, Blind Pen, Blind Cube, RGB Pen, RGB Cube (Complex Dynamics, Hard Exploration)

No tuning was needed for the actor-critic learning rate, the world model learning rates, and scaffolded / target policy learning rates.

In general, the paper focuses on the online, tabula rasa case where we are in an entirely unfamiliar environment. How adaptable is this method to either the offline case, or the finetuning case where we have an existing policy / world model?

Thank you for the interesting question. Indeed, we have thought about this too, and believe that Scaffolder may be well-suited for learning from offline data. Many large offline trajectory datasets, such as RT-X [1], have additional privileged information, in the form of metadata, captions, additional modalities, etc. that can be exploited by Scaffolder as privileged sensors to better train the target policy. Training Dreamer style models entirely offline may be challenging, but it may be possible to extend current offline model-based RL frameworks [2,3] to achieve this.

Scaffolder may also be an excellent fit for finetuning policies in settings where we have some prior knowledge. Scaffolder has explicit components to model dynamics models, policies, value functions, exploration policies etc. This permits injecting varied types of prior knowledge: for example, the privileged dynamics model could be initialized based on prior knowledge and finetuned. Residual learning methods [4], shown to have success in RL settings before, may also be combined with Scaffolder for further improved finetuning performance."

[1] Padalkar, Abhishek, et al. "Open x-embodiment: Robotic learning datasets and rt-x models." arXiv preprint arXiv:2310.08864 (2023).

[2] Kidambi, Rahul, et al. "Morel: Model-based offline reinforcement learning." Advances in neural information processing systems 33 (2020): 21810-21823.

[3] Yu, Tianhe, et al. "Mopo: Model-based offline policy optimization." Advances in Neural Information Processing Systems 33 (2020): 14129-14142.

[4] Haldar, Siddhant, et al. "Teach a Robot to FISH: Versatile Imitation from One Minute of Demonstrations." arXiv preprint arXiv:2303.01497 (2023).

Thank you for your review!

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.

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.

Related work section could be expanded to include research papers that leverage privileged simulation reset to improve policy.

That is an interesting point! While we mainly considered privileged training in the form of privileged sensors (hence the name sensory scaffolding), resets could be seen as another privileged training mechanism to scaffold policy learning. We have incorporated this discussion into the related work (changed text in red), citing the papers suggested.

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.

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.

While there has been considerable research in this area, as detailed in the related work, this paper adds value to the existing body of knowledge, even without introducing novel methods. The proposed method may not be groundbreaking, but it offers a comprehensive examination of this issue from four perspectives: model, value, representation, and exploration.

While this may be subjective, we respectfully disagree with the assertion that this paper does not introduce novel methods. Scaffolder is a novel method that permits privileged information to influence and improve policy learning through a comprehensive suite of routes: world model, value, exploration, and representation, as acknowledged by reviewers V3Vw and Sbzi. Firstly, it is not a priori clear how to achieve this, but furthermore, Scaffolder is more substantial than simply combining various existing methods. A particular salient example is Nested Latent Imagination, which frugally achieves the goals of generating synthetic experience for an impoverished target policy, from a world model powered by privileged sensory information. Note that others have also considered this problem and arrived at different solutions: see our baseline Informed Dreamer, which Scaffolder outperforms comprehensively.

Finally, beyond the method, this paper is also novel in considering learning from privileged information beyond just simulator state: Scaffolder can use high dimensional and noisy privileged observations like RGB images to train target policies.

Thank you for your review!

For scaffolded TD error comparison, it’s not clear why the comparison is conducted on Blind pick environment, since the gap between the proposed method and the version without scaffolded critic is much larger (at least in terms of relative gap) on Blind Cube Rotation environment. Also it would be great to see whether the estimate is close for tasks like Blind Locomotion (since the gap is small on that task).

Thank you, we believe this comment arises from a misunderstanding of the analysis experiment reported in Fig 9. Our exposition here grew rather terse for space, and we have updated the experiment description in Section 4.3 to improve it, following the clarifications below.

Fig 9 does not represent a comparison between Scaffolder and No Scaffolder Critic. Instead, it compares the policy return estimated with all scaffolded components (specifically dynamics, reward, critic, and continue heads, as seen in Eq 3) with the policy return estimated using no scaffolded components. Recall that the gradient of this estimated policy return determines the policy updates.

We ran this analysis on Blind Pick because, as Fig 8 shows, “No-Scaff. WM” performs very poorly here compared to full Scaffolder (this is also true for Blind Locomotion). We wanted to trace this performance gap all the way back to differences in the computed policy returns when not using scaffolded WM (including dynamics, reward, continue heads) and critic.

We also agree with the reviewer that reporting these results for tasks other than Blind Pick may be interesting, and plan to report them, hopefully within the rebuttal period.

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.

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), 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."

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.

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

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?

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.

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.

We have now performed the TD($\lambda$) analysis experiment on another environment, Blind Locomotion, as you had suggested. We find similar trends where the Scaffolded TD($\lambda$) return estimate is more accurate than the target TD($\lambda$) estimate. This provides further evidence that the training signals from Scaffolder are more accurate and lead to better policy improvement. Note that this experiment requires training Scaffolder from scratch with additional components (namely the target critic) to permit computing the target TD($\lambda$) estimate in addition to the scaffolded TD($\lambda$) estimate, so we had reported it only for Blind Pick (Fig 9) in the main paper. We added this Blind Locomotion result to Appendix D as additional evidence, alongside the Gridworld results already reported there.

We have provided some walltime learning curves, in response to xL3j, that you may also be interested in. Due to time constraints, we didn’t have time to rerun all methods on a single machine, but we will do so after the rebuttal.

Thank you for your review!

Sim-to-real scaffolder for real-world applications?

Yes, Scaffolder can easily be used as a sim-to-real approach, where simulator state serves as the privileged observation and sensory inputs (e.g. camera images) can serve as the target observation. Among our experimental settings, RGB Cube and RGB Pen involving privileged object pose and target visual observations follows a similar template.

Furthermore, note that we evaluate Scaffolder in more general settings that suggest a more direct route to real-world applications: even the privileged observations in our settings often come from additional sensors, rather than directly reading the simulator state. (e.g. privileged cameras in the Blind Pick environment) — this can be replicated in real-world learning settings such as the autonomous car setup described in the introduction. Scaffolder is also sample-efficient enough to reasonably permit this: it inherits and improves the sample-efficiency of Dreamer, which has already been applied to real-world robot learning [1].

Finally, you may also be interested in our response to reviewer Sbzi, where we suggest some practical ways to fine-tune and use offline data for Scaffolder.

[1] Wu, Philipp, et al. "Daydreamer: World models for physical robot learning." Conference on Robot Learning. PMLR, 2023.

Performance vs training wall-clock time?

Appendix B in the submission already reports the total training time for Scaffolder and baselines, showing that Scaffolder trains in similar time to other baselines. The model-free baselines RMA and AAC both need much longer to train than the model-based approaches like Scaffolder. Scaffolder needs on average 30-40% more wall clock time per update step than base Dreamer v3, due to having to train an additional world model but generally makes up for this by training much more sample-efficiently. We will expand this section to provide plots like Fig 6 but with wall clock time on the x-axis.

Finally, note that training time costs in most real-world applications would be dominated by the cost of collecting environment steps (“frames”), such as with a robot moving in the real-world at each step. As such, the number of environment steps is often the most relevant cost to consider, which is why research in this area usually reports performance vs. steps as in our Fig 6.

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^-$.

We have now completed all three pending experiments requested by reviewers:

• We evaluated Scaffolder's efficacy over SOTA baselines on a pre-existing privileged RL benchmark, as requested by Reviewer xL3j
• We verified that Scaffolder yields more accurate learning signals on one more environment, Blind Locomotion, as requested by Reviewer V3Vw.
• We plotted learning curves over walltime rather than environment samples and find that trends largely hold, as requested by Reviewer xL3j and V3Vw.

The details of these experiments are reported in the individual responses to the reviewers. Overall, these new results add further support to the claims in the paper. We would like to express our sincere gratitude to all reviewers for the constructive feedback and for helping to improve our paper.