Title: MoVie: Visual Model-Based Policy Adaptation for View Generalization
Venue: NeurIPS 2023
Accepted: yes
Keywords: visual reinforcement learning

visual generalization
Abstract: Visual Reinforcement Learning (RL) agents trained on limited views face significant challenges in generalizing their learned abilities to unseen views. This inherent difficulty is known as the problem of $\textit{view generalization}$. In this work, we systematically categorize this fundamental problem into four distinct and highly challenging scenarios that closely resemble real-world situations. Subsequently, we propose a straightforward yet effective approach to enable successful adaptation of visual $\textbf{Mo}$del-based policies for $\textbf{Vie}$w generalization ($\textbf{MoVie}$) during test time, without any need for explicit reward signals and any modification during training time. Our method demonstrates substantial advancements across all four scenarios encompassing a total of $\textbf{18}$ tasks sourced from DMControl, xArm, and Adroit, with a relative improvement of $\mathbf{33}$%, $\mathbf{86}$%, and $\mathbf{152}$% respectively. The superior results highlight the immense potential of our approach for real-world robotics applications. Code and videos are available at https://yangsizhe.github.io/MoVie/.

MoVie: Visual Model-Based Policy Adaptation for View Generalization

Abstract

Visual Reinforcement Learning (RL) agents trained on limited views face significant challenges in generalizing their learned abilities to unseen views. This inherent difficulty is known as the problem of view generalization . In this work, we systematically categorize this fundamental problem into four distinct and highly challenging scenarios that closely resemble real-world situations. Subsequently, we propose a straightforward yet effective approach to enable successful adaptation of visual Mo del-based policies for Vie w generalization ( MoVie ) during test time, without any need for explicit reward signals and any modification during training time. Our method demonstrates substantial advancements across all four scenarios encompassing a total of 18 tasks sourced from DMControl, xArm, and Adroit, with a relative improvement of 33 % , 86 % , and 152 % respectively. The superior results highlight the immense potential of our approach for real-world robotics applications. Code and videos are available at yangsizhe.github.io/MoVie.

0

200

400

600

493

521

427

656

DMControl

0

20

40

28

40

32

52

xArm

0

20

40

21

30

21

53

Adroit

TD-MPC

DM

IDM+STN

MoVie

1 Introduction

Visual Reinforcement Learning (RL) has achieved great success in various applications such as video games [17,18], robotic manipulation [22], and robotic locomotion [34]. However, one significant challenge for real-world deployment of visual RL agents remains: a policy trained with very limited views (commonly one single fixed view) might not generalize to unseen views. This challenge is especially pronounced in robotics, where a few fixed views may not adequately capture the variability of the environment. For instance, the RoboNet dataset [3] provides diverse views across a range of manipulation tasks, but training on such large-scale data only yields a moderate success rate ( 10% ∼ 20% ) for unseen views [3].

Recent efforts have focused on improving the visual generalization of RL agents [1,7,9,36]. However, these efforts have mainly concentrated on generalizing to different appearances and backgrounds. In contrast, view generalization presents a unique challenge as the deployment view is unknown and may move freely in the 3D space. This might weaken the weapons such as data augmentations [12, 14,15,33] that are widely used in appearance generalization methods. Meanwhile, scaling up domain

Training Time

TestTime

(Seen View)

(PolicyAdapationonUnseenView)

SpatialAdaptive

Ot

Encoder

LatentDynamics

Ot

Encoder

LatentDynamics

h

p

h

装

p

n

LMSE

LMSE

Action

TransformParam

Action

h

h

Concatenation

Encoder

=2D Transformation

SpatialAdaptive

Encoder

randomization [19,21,31] to all possible views is usually unrealistic because of the large cost and the offline nature of existing robot data. With these perspectives combined, it is difficult to apply those common approaches to address view generalization.

In this work, we commence by explicitly formulating the test-time view generalization problem into four challenging settings: a) novel view , where the camera changes to a fixed novel position and orientation; b) moving view , where the camera moves continuously around the scene, c) shaking view , where the camera experiences constant shaking, and d) novel FOV , where the field of view (FOV) of the camera is altered initially. These settings cover a wide spectrum of scenarios when visual RL agents are deployed to the real world. By introducing these formulations, we aim to advance research in addressing view generalization challenges and facilitate the utilization of robot datasets [3] and deployment on physical robotic platforms.

To address the view generalization problem, we argue that adaptation to novel views during test time is crucial, rather than aiming for view-invariant policies. We propose MoVie, a simple yet effective method for adapting visual Mo del-based policies to generalize to unseen Vie ws. MoVie leverages collected transitions from interactions and incorporates spatial transformer networks (STN [13]) in shallow layers, using the learning objective of the dynamics model (DM). Notably, MoVie requires no modifications during training and is compatible with various visual model-based RL algorithms. It only necessitates small-scale interactions for adaptation to the deployment view.

We perform extensive experiments on 7 robotic manipulation tasks (Adroit hand [20] and xArm [7]) and 11 locomotion tasks (DMControl suite [30])), across the proposed 4 view generalization settings, totaling 18 × 4 configurations. MoVie improves the view generalization ability substantially, compared to strong baselines including the inverse dynamics model (IDM [7]) and the dynamics model (DM). Remarkably, MoVie attains a relative improvement of 86% in xArm and 152% in Adroit, underscoring the potential of our method in robotics. We are committed to releasing our code and testing platforms. To conclude, our contributions are three-fold:

We formulate the problem of view generalization in visual reinforcement learning with a wide range of tasks and settings that mimic real-world scenarios.

We propose a simple model-based policy adaptation method for view generalization ( MoVie ), which incorporates STN into the shallow layers of the visual representation with a self-supervised dynamics prediction objective.

We successfully showcase the effectiveness of our method through extensive experiments. The results serve as a testament to its capability and underscore its potential for practical deployment in robotic systems, particularly with complex camera views.

2 Related Work

Visual generalization in reinforcement learning. Agents trained by reinforcement learning (RL) from visual observations are prone to overfitting the training scenes, making it hard to generalize to

unseen environments with appearance differences. A large corpus of recent works has focused on addressing this issue [7, 9, 10, 14, 15, 19, 21, 31, 33, 35, 36]. Notably, SODA [10] provides a visual generalization benchmark to better evaluate the generalizability of policies, while they only consider appearance changes of agents and backgrounds. Distracting control suite [26] adds both appearance changes and camera view changes into DMControl [30], where the task diversity is limited.

View generalization in robotics. The field of robot learning has long grappled with the challenge of training models on limited views and achieving generalization to unseen views. Previous studies, such as RoboNet [3], have collected extensive video data encompassing various manipulation tasks. However, even with pre-training on such large-scale datasets, success rates on unseen views have only reached approximately 10% ∼ 20% [3]. In recent efforts to tackle this challenge, researchers have primarily focused on third-person imitation learning [23-25] and view-invariant visual representations [2, 4, 16, 32, 37], but these approaches are constrained by the number of available camera views. In contrast, our work addresses a more demanding scenario where agents trained on a single fixed view are expected to generalize to diverse unseen views and dynamic camera settings.

Test-time training. There is a line of works that train neural networks at test-time with self-supervised learning in computer vision [5, 28, 29], robotics [27], and visual RL [7]. Specifically, PAD is the closest to our work [7], which adds an inverse dynamics model (IDM) objective into model-free policies for both training time and test time and gains better appearance generalization. In contrast, we differ in a lot of aspects: (i) we focus on visual model-based policies, (ii) we require no modification in training time, and (iii) our method is designed for view generalization specifically.

3 Preliminaries

Formulation. We model the problem as a Partially Observable Markov Decision Process (POMDP) M = ⟨O , A , T , R , γ ⟩ , where o ∈ O are observations, a ∈ A are actions, F : O × A ↦→ O is a transition function (called dynamics as well), r ∈ R are rewards, and γ ∈ [0 , 1) is a discount factor. During training time, the agent's goal is to learn a policy π that maximizes discounted cumulative rewards on M , i.e. , max E π [ ∑ ∞ t =0 γ t r t ] . During test time, the reward signal from the environment is not accessible to agents and only observations are available, which are possible to experience subtle changes such as appearance changes and camera view changes.

Model-based reinforcement learning. TD-MPC [11] is a model-based RL algorithm that combines model predictive control and temporal difference learning. TD-MPC learns a visual representation z = h ( o ) that maps the high-dimensional observation o ∈ O into a latent state z ∈ Z and a latent dynamics model d : Z × A ↦→ Z that predicts the future latent state z ′ = d ( z , a ) based on the current latent state z and the action a . MoDem [8] accelerates TD-MPC with efficient utilization of expert demonstrations D = { D 1 , D 2 , · · · , D N } to solve challenging tasks such as dexterous manipulation [20]. In this work, we select TD-MPC and MoDem as the backbone algorithm to train model-based agents, while our algorithm could be easily extended to most model-based RL algorithms such as Dreamer [6].

4 Method

We propose a simple yet effective method, visual Mo del-based policy adaptation for Vie w generalization ( MoVie ), which can accommodate visual RL agents to novel camera views at test time.

Learning objective for the test time. Given a tuple ( o t , a t , o t +1 ) , the original latent state dynamics prediction objective can be written as

where h is an image encoder that projects a high-dimensional observation from space O into a latent space Z and d is a latent dynamics model d : Z × A ↦→ Z .

In test time, the observations under unseen views lie in a different space O ′ , so that their corresponding latent space also changes to Z ′ . However, the projection h learned in training time can only map O ↦→ Z while the policy π only learns the mapping Z ↦→ A , thus making the policy hard to generalize to the correct mapping function Z ′ ↦→A . Our proposal is thus to adapt the projection h from a mapping function h : O ↦→ Z to a more useful mapping function h ′ : O ′ ↦→Z so that the

Baseline

Ours

Z

Z

A

Training

h

π

h

Time

一

D

p

Z'

Test

h

ⅣX

hSAE

Time

TransformParam

TransformParam

Localisation

Localisation

Net

Net

conv

ReLU

conv

ReLU

Zt

中

GridSampler

GridSampler

Ot

MLP

policy would execute the correct mapping Z ↦→ A without training. A vivid illustration is provided in Figure 3.

We freeze the latent dynamics model d , denoted as d ⋆ , so that the latent dynamics model is not a training target but a supervision. We also insert STN blocks [13] into the shallow layers of h to better adapt the projection h , so that we write h as h SAE (SAE denotes spatial adaptive encoder). Though the objective is still the latent state dynamics prediction loss, the supervision here is superficially identical but fundamentally different from training time. The formal objective is written as

Spatial adaptive encoder. We now describe more details about our modified encoder architecture during test time, referred to as spatial adaptive encoder (SAE). To keep our method simple and fast to adapt, we only insert two different STNs into the original encoder, as shown in Figure 4. We observe in our experiments that transforming the low-level features ( i.e. , RGB features and shallow layer features) is most critical for adaptation, while the benefit of adding more STNs is limited (see Table 6). An STN block consists of two parts: (i) a localisation net that predicts an affine transformation with 6 parameters and (ii) a grid sampler that generates an affined grid and samples features from the original feature map. The point-wise affine transformation is written as

where G is the sampling grid, ( x t i , y t i ) are the target coordinates of the regular grid in the output feature map, ( x s i , y s i ) are the source coordinates in the input feature map that define the sample points, and A ϕ is the affine transformation matrix.

Training strategy for SAE. We use a learning rate 1 × 10 -5 for STN layers and 1 × 10 -7 for the encoder. We utilize a replay buffer with size 256 to store history observations and update 32 times for each time step to heavily utilize the online data. Implementation details remain in Appendix A.

Training View

Novel View

Moving View

Shaking View

Novel FOV

5 Experiments

In this section, we investigate how well an agent trained on a single fixed view generalizes to unseen views during test time. During the evaluation, agents have no access to reward signals, presenting a significant challenge for agents to self-supervise using online data.

5.1 Experiment Setup

Formulation of camera view variations : a) novel view , where we maintain a fixed camera target while adjusting the camera position in both horizontal and vertical directions by a certain margin, b) moving view , where we establish a predefined trajectory encompassing the scene and the camera follows this trajectory, moving back and forth for each time step while focusing on the center of the scene, c) shaking view , where we add Gaussian noise onto the original camera position at each time step, and d) novel FOV , where the FOV of the camera is altered once, different from the training phase. A visualization of four settings is provided in Figure 6 and videos are available in yangsizhe.github.io/MoVie for a better understanding of our configurations. Details remain in Appendix B.

Tasks. Our test platform consists of 18 tasks from 3 domains: 11 tasks from DMControl [30], 3 tasks from Adroit [20], and 4 tasks from xArm [7]. A visualization of the three domains is in Figure 5. Although the primary motivation for this study is for addressing view generalization in real-world robot learning, which has not yet been conducted, we contend that the extensive range of tasks tackled in our research effectively illustrates the potential of MoVie for real-world application. We run 3 seeds per experiment with seed numbers 0 , 1 , 2 and run 20 episodes per seed. During these 20 episodes, the models could store the history transitions. We report cumulative rewards for DMControl tasks and success rates for robotic manipulation tasks, averaging over episodes.

Baselines. We first train visual model-based policies with TD-MPC [11] on DMControl and xArm environments and MoDem [8] on Adroit environments under the default configurations and then test these policies in the view generalization setting. Since MoVie consists of two components mainly ( i.e. , DMand STN), we build two test-time adaptation algorithms by replacing each module: a) DM , which removes STN from MoVie and b) IDM+STN , which replaces DM in MoVie with IDM [7]. IDM+STN is very close to PAD [7], and we add STN for fair comparison. We keep all training settings the same.

5.2 Main Experiment Results

Considering the extensive scale of our conducted experiments, we present an overview of our findings in Table 1 and Figure 1. The detailed results for four configurations are provided in Table 2, Table 3, Table 4, and Table 5 respectively. We then detail our findings below.

Superiority of MoVie across domains, especially in robotic manipulation. It is observed that irrespective of the domain or configuration, MoVie consistently outperforms TD-MPC without any adaptation and other methods that apply DM and IDM. This large improvement highlights the fundamental role of our straightforward approach in addressing view generalization. Additionally, we observe that MoVie exhibits greater suitability for robotic manipulation tasks. We observe a significant relative improvement of 86 % in xArm tasks and 152 % in Adroit tasks, as opposed to a comparatively modest improvement of only 33% in DMControl tasks; see Table 1. We attribute this

disparity to two factors: a) the inherent complexity of robotic manipulation tasks and b) the pressing need for effective view generalization in this domain.

Challenges in handling shaking view. Despite improvements of MoVie in various settings, we have identified a relative weakness in addressing the shaking view scenario. For instance, in xArm tasks, the success rate of MoVie is only 45% , which is close to the 42% success rate of TD-MPC without adaptation. Other baselines such as DM and IDM+STN also experience performance drops. We acknowledge the inherent difficulty of the shaking view scenario, while it is worth noting that in real-world robotic manipulation applications, cameras often exhibit smoother movements or are positioned in fixed views, partially mitigating the impact of shaking view.

Effective adaptation in novel view, moving view, and novel FOV scenarios. In addition to the shaking view setting, MoVie consistently outperforms TD-MPC without adaptation by 2 × ∼ 4 × in robotic manipulation tasks across the other three settings. It is worth noting that these three settings are among the most common scenarios encountered in real-world applications.

Real-world implications. Our findings have important implications for real-world deployment of robots. Previous methods, relying on domain randomization and extensive data augmentation during training, often hinder the learning process. Our proposed method enables direct deployment of offline or simulation-trained agents, improving success rates with minimal interactions.

5.3 Ablations

To validate the rationale behind several choices of MoVie and our test platform, we perform a comprehensive set of ablation experiments.

Integration of STN with low-level features. To enhance view generalization, we incorporate two STN blocks [13], following the image observation and the initial convolutional layer of the feature encoder. This integration is intended to align the low-level features with the training view, thereby preserving the similarity of deep semantic features to the training view. As shown in Table 6, by progressively adding more layers to the feature encoder, we observe that deeper layers do not provide significant additional benefits, supporting our intuition for view generalization.

Different novel views. We classify novel views into three levels of difficulty, with our main experiments employing the medium difficulty level by default. Table 7 presents additional results for

the easy and hard difficulty levels. As the difficulty level increases, we observe a consistent decrease in the performance of all the methods.

The efficacy of STN in conjunction with IDM. Curious readers might be concerned about our direct utilization of IDM+STN in the main experiments, suggesting that STN could potentially be detrimental to IDM. However, Table 8 shows that STN not only benefits our method but also improves the performance of IDM, thereby demonstrating effectiveness of SAE for adaptation of the representation and validating our baseline selection.

Finetune or freeze the DM. In our approach, we employ the frozen DM as a form of supervision to guide the adaptation process of the encoder. However, it remains unverified for the readers whether end-to-end finetuning of both the DM and the encoder would yield similar benefits. The results presented in Table 9 demonstrate that simplistic end-to-end finetuning does not outperform MoVie, thereby reinforcing the positive results obtained by MoVie.

6 Conclusion

In this study, we present MoVie, a method for adapting visual model-based policies to achieve view generalization. MoVie mainly finetunes a spatial adaptive image encoder using the objective of the latent state dynamics model during test time. Notably, we maintain the dynamics model in a frozen state, allowing it to function as a form of self-supervision. Furthermore, we categorize the view generalization problem into four distinct settings: novel view, moving view, shaking view, and novel FOV. Through a systematic evaluation of MoVie on 18 tasks across these four settings, totaling 64 different configurations, we demonstrate its general and remarkable effectiveness.

One limitation of our work is the lack of real robot experiments, while our focus is on addressing view generalization in robot datasets and deploying visual reinforcement learning agents in real-world scenarios. In our future work, we would evaluate MoVie on real-world robotic tasks.

Acknowledgment

This work is supported by National Key R&D Program of China (2022ZD0161700).

References

David Bertoin, Adil Zouitine, Mehdi Zouitine, and Emmanuel Rachelson. Look where you look! saliency-guided q-networks for generalization in visual reinforcement learning. NeurIPS , 2022. 1, 17

Boyuan Chen, Pieter Abbeel, and Deepak Pathak. Unsupervised learning of visual 3d keypoints for control. In ICML , 2021. 3

Sudeep Dasari, Frederik Ebert, Stephen Tian, Suraj Nair, Bernadette Bucher, Karl Schmeckpeper, Siddharth Singh, Sergey Levine, and Chelsea Finn. Robonet: Large-scale multi-robot learning. CoRL , 2019. 1, 2, 3

Danny Driess, Ingmar Schubert, Pete Florence, Yunzhu Li, and Marc Toussaint. Reinforcement learning with neural radiance fields. arXiv , 2022. 3

Yossi Gandelsman, Yu Sun, Xinlei Chen, and Alexei Efros. Test-time training with masked autoencoders. NeurIPS , 2022. 3

Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv , 2020. 3

Nicklas Hansen, Rishabh Jangir, Yu Sun, Guillem Alenyà, Pieter Abbeel, Alexei A Efros, Lerrel Pinto, and Xiaolong Wang. Self-supervised policy adaptation during deployment. ICLR , 2021. 1, 2, 3, 5, 13, 17

Nicklas Hansen, Yixin Lin, Hao Su, Xiaolong Wang, Vikash Kumar, and Aravind Rajeswaran. Modem: Accelerating visual model-based reinforcement learning with demonstrations. ICLR , 2023. 3, 5, 13

Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. NeurIPS , 2021. 1, 3, 17

Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation. In International Conference on Robotics and Automation , 2021. 3

Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. ICML , 2022. 3, 5, 13

Nicklas Hansen, Zhecheng Yuan, Yanjie Ze, Tongzhou Mu, Aravind Rajeswaran, Hao Su, Huazhe Xu, and Xiaolong Wang. On pre-training for visuo-motor control: Revisiting a learningfrom-scratch baseline. ICML , 2023. 1

Max Jaderberg, Karen Simonyan, Andrew Zisserman, et al. Spatial transformer networks. NeurIPS , 2015. 2, 4, 8

Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. NeurIPS , 2020. 1, 3

Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. arXiv preprint arXiv:1910.05396 , 2019. 1, 3

Yunzhu Li, Shuang Li, Vincent Sitzmann, Pulkit Agrawal, and Antonio Torralba. 3d neural scene representations for visuomotor control. In CoRL , 2022. 3

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv , 2013. 1

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature , 2015. 1

Xue Bin Peng, Marcin Andrychowicz, Wojciech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In ICRA , 2018. 2, 3

Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. RSS , 2018. 2, 3, 5, 13

Fabio Ramos, Rafael Carvalhaes Possas, and Dieter Fox. Bayessim: adaptive domain randomization via probabilistic inference for robotics simulators. arXiv , 2019. 2, 3

Rutav Shah and Vikash Kumar. Rrl: Resnet as representation for reinforcement learning. ICML , 2021. 1

Jinghuan Shang and Michael S Ryoo. Self-supervised disentangled representation learning for third-person imitation learning. In IROS , 2021. 3

Pratyusha Sharma, Deepak Pathak, and Abhinav Gupta. Third-person visual imitation learning via decoupled hierarchical controller. NeurIPS , 2019. 3

Bradly C Stadie, Pieter Abbeel, and Ilya Sutskever. Third-person imitation learning. ICLR , 2017. 3

Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite-a challenging benchmark for reinforcement learning from pixels. arXiv , 2021. 3

Yu Sun, Wyatt L Ubellacker, Wen-Loong Ma, Xiang Zhang, Changhao Wang, Noel V CsomayShanklin, Masayoshi Tomizuka, Koushil Sreenath, and Aaron D Ames. Online learning of unknown dynamics for model-based controllers in legged locomotion. RA-L , 2021. 3

Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In ICML , 2020. 3

Yu Sun, Xiaolong Wang, Liu Zhuang, John Miller, Moritz Hardt, and Alexei A. Efros. Test-time training with self-supervision for generalization under distribution shifts. In ICML , 2020. 3

Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv , 2018. 2, 3, 5, 13

Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In IROS , 2017. 2, 3

Hsiao-Yu Fish Tung, Zhou Xian, Mihir Prabhudesai, Shamit Lal, and Katerina Fragkiadaki. 3d-oes: Viewpoint-invariant object-factorized environment simulators. arXiv , 2020. 3

Kaixin Wang, Bingyi Kang, Jie Shao, and Jiashi Feng. Improving generalization in reinforcement learning with mixture regularization. NeurIPS , 2020. 1, 3

Ruihan Yang, Minghao Zhang, Nicklas Hansen, Huazhe Xu, and Xiaolong Wang. Learning vision-guided quadrupedal locomotion end-to-end with cross-modal transformers. ICLR , 2021. 1

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv , 2021. 3, 17

Zhecheng Yuan, Zhengrong Xue, Bo Yuan, Xueqian Wang, Yi Wu, Yang Gao, and Huazhe Xu. Pre-trained image encoder for generalizable visual reinforcement learning. NeurIPS , 2022. 1, 3, 17

Yanjie Ze, Nicklas Hansen, Yinbo Chen, Mohit Jain, and Xiaolong Wang. Visual reinforcement learning with self-supervised 3d representations. RA-L , 2023. 3

Appendix

A Implementation Details

In this section, we describe the implementation details of our algorithm for training on the training view and test time training in view generalization settings on the DMControl [30], xArm [7], and Adroit [20] environments. We utilize the official implementation of TD-MPC [11] and MoDem [8] which are available at github.com/nicklashansen/tdmpc and github.com/facebookresearch/modem as the model-based reinforcement learning codebase. During training time, we use the default hyperparameters in official implementation of TD-MPC and MoDem. We present relevant hyperparameters during both training and test time in Table 10 and Table 11. One seed of our experiments could be run on a single 3090 GPU with fewer than 2 GB and it takes ∼ 1 hours for test-time training.

Training time setup. We train visual model-based policies with TD-MPC on DMControl and xArm environments, and MoDem on Adroit environments, We employ identical network architecture and hyperparameters as original TD-MPC and MoDem during training time.

The network architecture of the encoder in original TD-MPC is composed of a stack of 4 convolutional layers, each with 32 filters, no padding, stride of 2, 7 × 7 kernels for the first one, 5 × 5 kernels for the second one and 3 × 3 kernels for all others, yielding a final feature map of dimension 3 × 3 × 32 (inputs whose framestack is 3 have dimension 84 × 84 × 9). After the convolutional layers, a fully connected layer with an input size of 288 performs a linear transformation on the input and generates a 50-dimensional vector as the final output.

The network architecture of the encoder in original Modem is composed of a stack of 6 convolutional layers, each with 32 filters, no padding, stride of 2, 7 × 7 kernels for the first one, 5 × 5 kernels for the second one and 3 × 3 kernels for all others, yielding a final feature map of dimension 2 × 2 × 32 (inputs whose framestack is 2 have dimension 224 × 224 × 6). After the convolutional layers, a fully connected layer with an input size of 128 performs a linear transformation on the input and generates a 50-dimensional vector as the final output.

Test time training setup. During test time, we train spatial adaptive encoder (SAE) to adapt to view changes. We insert STN blocks before and after the first convolutional layer of the original encoders in TD-MPC and MoDem. The original encoders are augmented by inserting STN blocks, resulting in the formation of SAE. Particularly, for the STN block inserted before the first convolutional layer, the input is a single frame. This means that when the frame stack size is N, N individual frames are fed into this STN block. This is done to apply different transformations to different frames in cases of moving and shaking view.

To update the SAE, we collect online data using a buffer with a size of 256. For each update, we randomly sample 32 (observation, action, next_observation) tuples from the buffer as a batch. The optimization objective is to minimize the loss in predicting the dynamics of the latent states, as defined in Equation 2.

During testing on each task, we run 20 consecutive episodes, although typically only a few or even less than one episode is needed for the test-time training to converge. To make efficient use of the data collected with minimal interactions, we employ a multi-update strategy. After each interaction with the environment, the SAE is updated 32 times.

The following is the network architecture of the first STN block inserted into the encoder of TD-MPC.

B Environment Details

We categorize the view generalization problem into four distinct settings: novel view, moving view, shaking view, and novel FOV. In this section, we provide descriptions of the implementation details for

each setting. The detailed camera settings can be referred to in the code of the environments that we are committed to releasing or in the visualization available on our website yangsizhe.github.io/MoVie.

Novel view. In this setting, for locomotion tasks (cheetah-run, walker-stand, walker-walk, and walker-run), the camera always faces the moving agent, while for other tasks, the camera always faces a fixed point in the environment. Therefore, as we change the camera position, the camera orientation also changes accordingly.

Moving view. Similar to the previous setting, the camera also always faces the moving agent or a fixed point in the environment. The camera position varies continuously.

Shaking view. To simulate camera shake, we applied Gaussian noise to the camera position (XYZ coordinates in meters) at each time step. For DMControl and Adroit, the mean of the distribution is 0, the standard deviation is 0.04, and we constrain the noise within the range of -0.07 to +0.07. For xArm, the mean of the distribution is 0, the standard deviation is 0.4, and we constrain the noise within the range of -0.07 to +0.07.

Novel FOV. We experiment with a larger FOV. For DMControl, we modify the FOV from 45 to 53. For xArm, we modify the FOV from 50 to 60. For Adroit, we modify the FOV from 45 to 50. We also experiment with a smaller FOV and results are presented in Appendix F.

C Visualization of Feature Maps from Shallow to Deep Layers

We incorporate STN in the first two layers of the visual encoder for all tasks. After visualizing the features of different layers (as shown in Figure 7), we found that the features of shallow layers contain

more information about the spatial relationship. Therefore transforming the features of shallow layers for view generalization is reasonable.

.

D Visualization of Feature Map Transformation

We visualize the first layer feature map of the image encoder from TD-MPC and MoVie in Figure 8. It is observed that the feature map from MoVie on the novel view exhibits a closer resemblance to that on the training view.

Training View

Novel View

TD-MPC

Ours

E Extended Description of Baselines

TD-MPC. Wetest the agent trained on training view without any adaptation in the view generalization settings.

DM. This is derived from MoVie by removing STN blocks, which just adapts encoder during test time.

IDM+STN. This is derived from MoVie by replacing the dynamics model with the inverse dynamics model which predicts the action in between based on the latent states before and after transition. The inverse dynamics model is finetuned together with the encoder and STN blocks during testing.

F Ablation on Different FOVs

In our main experiments, we consider the novel FOV as a FOV larger than the original. In Table 12, we present results for both smaller and larger FOV scenarios. Our method demonstrates the successful handling of both cases.

G Comparison with Other Visual RL Generalization Algorithms

In addition to baselines in the main paper, we also compare MoVie with state-of-the-art methods that focus on visual generalization or data augmentation in visual RL, i.e. , DrQ-v2 [35], SVEA [9], PIE-G [36], SGQN [1] and PAD [7]. Results on DMControl environments are shown in Table 13. And the training performance is shown in Table 14. It is observed that MoVie still outperforms these two methods that do not adapt in test time, while PIE-G could also achieve reasonable returns under the view perturbation. This may indicate that fusing the pre-trained networks from PIE-G into MoVie might lead to better results, which we leave as future work. Note that unlike MoVie, other methods need strong data augmentation or modification during training time.

H Results of Original Models on Training View

The performance of the original agents without any adaptation under the training view is reported in Table 15, 16, and 17 for reference. In the context of view generalization, it is evident that the performance of agents without adaptation significantly deteriorates.

Peer reviews:
comments: SUMMARY:
This work presents an approach to train model-based RL methods such that it generalize to novel views on multiple RL benchmarks. The method leverages classical STN and frozen encoder to benefit the generalization performance.

STRENGTHS:
The method is sound and simple with barely hyperparameters tuning

The design choices study such as different moving views and the insight are helpful for research community.

WEAKNESSES:
The writing can be largely improved on section 2 and 3. It is unclear what the model-based RL problem is, and how to define the view generalization, how to map the model to actual control/policy for evaluations, for audience that are not familiar with this line of work.

The problem is specific and the method novelty are limited.

Comparisons with baselines are unclear. At least some comparison with model-free RL and recent work [1] would be helpful.
[1] Multi-View Masked World Models for Visual Robotic Manipulation, Seo et al., 2023

QUESTIONS:
In the experiments, does the baseline methods see all various views in the training sets as well? I.e. is the baseline method in figure 3 also seeing different views?

There are many other choices to solve this problem. Is contrastive learning / or data augmentations on different views also a way to map different views to the same representations?

I am not sure the application on real-world robot experiments, although some visualizations are presented.

LIMITATIONS:
See above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-06-17

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper mainly provides a training paradigm. Utilizing the dynamic transition model of the environment during the testing phase as a supervisory signal, STN is used to quickly finetune the mapping of observed potential states, resulting in better performance of the strategy mapping trained on a single view task on unseen test views. This work conducted comparative experiments on three types of tasks and four challenging views, achieving results that surpass existing methods. Ablation experiments were conducted to prove the effectiveness of each module setting.\n\nSTRENGTHS:\nThe training paradigm provided in this article can effectively handle the generalization problem when view changes occur, and a thorough comparative experiment and ablation experiment have been conducted on the proposed model. The core idea of reconstructing mapping h using the environmental dynamic transfer model in the paper is effective.\n\nWEAKNESSES:\n1. The subscripts for o and a in Formula 2 are missing.\n2. There are issues with the baseline selection of IDM+STN. Cheetah run alone cannot prove that IDM+STN is generally better than IDM. It is not a problem to conduct ablation experiments on the Cheetah run model alone, but this cannot be used as sufficient evidence for the superiority of IDM+STN over IDM and thus selecting it as the baseline.\n3. Chapter 3 emphasizes the differences between Formula 2 and Formula 1, on one hand, due to the fixed parameters of Network d, and on the other hand due to the use of SAE in Network h. One of the core methods is to fix the parameters of d, which is intuitive to avoid changes in the target domain of h during the update process, resulting in the failure of the policy pi. However, the ablation experiment showed that the setting of d^* did not achieve a consistently effective effect.\n4. The main approach is to use the dynamic transfer model of the environment under the new view as supervision to quickly train the changed h '. Can you obtain a generalized h directly based on o'?\n\nQUESTIONS:\n1. Have the training methods of the three baselines in the comparative experiment also been finetuned? How much data was used during finetuning, and is this amount of data acceptable for practical applications?\n2. Can the network of h ^ SAE be directly used in training to achieve learning from multiple perspectives of tasks? Directly generalize from an unseen perspective without finetuning.\n3. The reason why Movie is better than IDM+STN should be explained.\n\nLIMITATIONS:\nThe authors have discussed the limitations", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses the novel problem of view generalization in reinforcement learning, where an RL agent is trained on an environment with a fixed view and then evaluated on a test environment having the exact same dynamics but observed from a different perspective. In order to address this issue, the authors introduce an innovative method that permits test-time adaptation to the new view. \nThe authors integrate a learnable spatial transformer network (STN) into the feature extraction component of the agent. This augmented encoder is then fine-tuned at test time to generate features that enable a frozen latent dynamic model to predict future state representations. An essential aspect of the authors' approach is that this adaptation does not necessitate any reward signal, which is a key advantage.\nExperimental results demonstrate a significant reduction in the generalization gap presented by a new view across several heterogeneous benchmarks.\n\nSTRENGTHS:\nThe paper is well-structured and easy to follow, providing clear and detailed explanations that make  believe that I could readily reproduce the method based on the descriptions provided.  \nThe authors introduce a new problem in the field that could stimulate further research.  \nThe proposed method's strength lies in its test-time adaptation capability, which eliminates the need for any reward signal, making it highly applicable to real-world scenarios. This characteristic also suggests that the approach could be integrated into any model-based method, enhancing its universality.  \nAnother advantage is the apparent minimal interaction required with the testing environment, promoting efficiency.  \nThe method is evaluated across various environments, with the results showcasing impressive improvements compared to non-adaptive methods.  \nThe paper's ablation study effectively illustrates the contribution of each component of the method, highlighting the importance of every aspect in achieving the observed results.  Additionally, the annex offers excellent visualizations demonstrating the impact of the spatial transformer on the feature map, providing intuitive understanding of the method's mechanics.\n\nWEAKNESSES:\nDespite the many strengths of the paper, a few areas could benefit from further development and clarification.\n\n1. It would be beneficial to see a comparison of the performance decrease relative to the original view. The lack of this data makes it challenging to truly gauge the significance of the generalization gap and the efficacy of the proposed method.\n\n2. The scope of comparative studies is somewhat limited. The paper primarily compares with methods not designed for visual adaptation, which may not provide the most insightful comparison. It would have been advantageous to see how the proposed approach stacks up against other strategies specifically intended for visual adaptation.\n\n3. An exploration of whether the proposed method can also benefit model-free algorithms such as Soft Actor-Critic (SAC) is missing. This could broaden the applicability of the findings and provide additional insights into the method's versatility.\n\n4. It remains unclear why the proposed method underperforms with the Inverse Dynamics Model (IDM). A deeper exploration of this anomaly could bolster the robustness and reliability of the approach.\n\n5. Finally, the similarity between the proposed method and the PAD method, which performs test-time domain adaptation in model-free RL, raises questions. The key differences are the use of a Spatial Transformer Network and the leveraging of the latent dynamics model of model-based methods instead of adding it as an auxiliary component. However, these differences, while noteworthy, do not necessarily represent a substantial departure from the PAD method. This similarity begs the question of whether the presented method offers significant novelty or if it is essentially an adjustment of existing approaches.\n\nQUESTIONS:\n1. To better understand the importance of the generalization gap, could you provide performance data related to the original view? This information would help quantify the performance decrease when the agent is switched to a different view.\n\n2. The use of inverse dynamics models with spatial transformer networks in your method isn't entirely clear. \nDoes IDM completely replace the DM in the TD-MPC algorithm?\nOr does it just used as a side network during TD-MPC training and used for test adaptation only?\n\nLIMITATIONS:\nThe authors mentioned a limitation of this work: the proposed method has not been tested in real-world scenarios, such as with physical robots, which may limit its immediate applicability.\nAnother potential limitation could be that the method has only been tested on model-based reinforcement learning (RL), whereas it would be relatively straightforward to evaluate its effectiveness in model-free RL.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on improving the generalization ability of visual DRL to adapt to unseen views. The authors propose a model-based policy adaptation approach that combines spatial transformer networks with a self-supervised dynamics prediction objective to address this problem. The effectiveness of the approach is evaluated through experiments on three commonly used benchmarks.\n\nSTRENGTHS:\n- This paper is well-written and easy to follow. \n- The proposed method is well-motivated. \n- The authors try to solve a realistic and essential problem, i.e., view generalization, especially in real-world or sim-to-real robotic scenarios.\n\nWEAKNESSES:\n- While the problem setting is attractive, the technical contribution is insufficient as a full NeurIPS paper. It is a straightforward combination of multiple existing works, like STN. \n- There needs to be more than the current evaluation to demonstrate the superiority over existing works.\n\nQUESTIONS:\n- The experiments need to be more comprehensive. The comparison baselines used are more like ablation studies rather than state-of-the-art (SOTA) methods, such as TD-MPC (the backbone of MoVie), DM (MoVie without STN), and IDM+STN (MoVie replacing DM with IDM). The authors should consider incorporating more commonly-used or SOTA data augmentation algorithms for visual DRL, such as DrQ-v2[1] and SEVA[2].\n- The authors mention multiple times that only shallow STNs can improve performance, but they fail to provide deeper analysis, such as visual illustrations. Is the number of layers set for all tasks, or do different tasks require their own suitable layers? This information would be valuable.\n- The authors claim that \"Our proposed method enables direct deployment of offline or simulation-trained agents ...\", However, the lack of practical experimental demonstrations, such as robot manipulation, makes their claim unconvincing. I strongly suggest that the authors conduct real-world robot experiments to assess the applicability of their proposed method.\n\nReferences:\n- [1] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645, 2021.\n- [2] Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision 362 transformers under data augmentation. Advances in Neural Information Processing Systems, 34, 2021.\n\nLIMITATIONS:\nNot Applicable.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

comments: SUMMARY:
Adapting policy into new view setup is an important task in RL. This paper presents MoVie, Model-based policies for View generalization to achieve the fast view adaptation of model-based policy. Specifically, they combine spatial transformer networks into the encoder models and train it during test-time using dynamics prediction loss. Despite the simplicity, the experimental results show that the strong performance of MoVie over several benchmarks on various test set.

STRENGTHS:
1. The proposed method is simple yet provide stable performance gain over various methods. Experimental validation is extensive enough. 
2. The paper is well written and easy to follow.

WEAKNESSES:
1. Technical novelty and technical depth is limited. I might decrease scores if I found works that combine STN with test-time training. 

Besides, few discussions are made on the choice of test-time loss and entire architecture. For example, the experimental results suggest that the DM is better than IDM especially when combining test-time training, but the paper lacks discussion on why it is. Besides, I agree prediction loss is one of the natural loss to train STN, but I think it is also possible to extend the method to model-free methods by using some self-supervised loss function (like contrastive learning). 

2. I think the experimental validation is already comprehensive, but it lacks qualitative analysis of their representation before/after adaptation. 

3. No theoretical justification for the choice of test-time loss and or entire architecture.

QUESTIONS:
See weakness section.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09