SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sizhe Yang $^{12*}$ , Yanjie Ze $^{13*}$ , Huazhe Xu $^{415}$
$^{1}$ Shanghai Qi Zhi Institute,  $^{2}$ University of Electronic Science and Technology of China,
$^{3}$ Shanghai Jiao Tong University,  $^{4}$ Tsinghua University,  $^{5}$ Shanghai AI Lab
yangsizhe.github.io/MoVie
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 Model-based policies for View 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.
Figure 1: View generalization results across 3 domains. Our method MoVie largely improves the test-time view generalization ability, especially in robotic manipulation domains.
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\% \sim 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
Figure 2: Overview of MoVie. During training (left), the agent is trained with the latent dynamics loss. At test time (right), we freeze the dynamics model and modify the encoder as a spatial adaptive encoder to adapt the agents to novel views.
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 Model-based policies to generalize to unseen Views. 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 \times 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.
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\% \sim 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.
**Formulation.** We model the problem as a Partially Observable Markov Decision Process (POMDP)  $\mathcal{M} = \langle \mathcal{O}, \mathcal{A}, \mathcal{T}, \mathcal{R}, \gamma \rangle$ , where  $\mathbf{o} \in \mathcal{O}$  are observations,  $\mathbf{a} \in \mathcal{A}$  are actions,  $\mathcal{F}: \mathcal{O} \times \mathcal{A} \mapsto \mathcal{O}$  is a transition function (called dynamics as well),  $r \in \mathcal{R}$  are rewards, and  $\gamma \in [0,1)$  is a discount factor. During training time, the agent's goal is to learn a policy  $\pi$  that maximizes discounted cumulative rewards on  $\mathcal{M}$ , i.e.,  $\max_{\pi} [\sum_{t=0}^{\infty} \gamma^{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  $\mathbf{z} = h(\mathbf{o})$  that maps the high-dimensional observation  $\mathbf{o} \in \mathcal{O}$  into a latent state  $\mathbf{z} \in \mathcal{Z}$  and a latent dynamics model  $d: \mathcal{Z} \times \mathcal{A} \mapsto \mathcal{Z}$  that predicts the future latent state  $\mathbf{z}' = d(\mathbf{z}, \mathbf{a})$  based on the current latent state  $\mathbf{z}$  and the action  $\mathbf{a}$ . MoDem [8] accelerates TD-MPC with efficient utilization of expert demonstrations  $\mathcal{D} = \{D_1, D_2, \dots, 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].
We propose a simple yet effective method, visual Model-based policy adaptation for View generalization (MoVie), which can accommodate visual RL agents to novel camera views at test time.
Learning objective for the test time. Given a tuple  $(\mathbf{o}_t, \mathbf{a}_t, \mathbf{o}_{t+1})$ , the original latent state dynamics prediction objective can be written as
$$
\mathcal {L} _ {\text {d y n a m i c s}} = \left\| d \left(h \left(\mathbf {o} _ {t}\right), \mathbf {a} _ {t}\right) - h \left(\mathbf {o} _ {\mathbf {t} + \mathbf {1}}\right) \right\| _ {2}, \tag {1}
$$
where  $h$  is an image encoder that projects a high-dimensional observation from space  $\mathcal{O}$  into a latent space  $\mathcal{Z}$  and  $d$  is a latent dynamics model  $d: \mathcal{Z} \times \mathcal{A} \mapsto \mathcal{Z}$ .
In test time, the observations under unseen views lie in a different space  $\mathcal{O}'$ , so that their corresponding latent space also changes to  $\mathcal{Z}'$ . However, the projection  $h$  learned in training time can only map  $\mathcal{O} \mapsto \mathcal{Z}$  while the policy  $\pi$  only learns the mapping  $\mathcal{Z} \mapsto \mathcal{A}$ , thus making the policy hard to generalize to the correct mapping function  $\mathcal{Z}' \mapsto \mathcal{A}$ . Our proposal is thus to adapt the projection  $h$  from a mapping function  $h: \mathcal{O} \mapsto \mathcal{Z}$  to a more useful mapping function  $h': \mathcal{O}' \mapsto \mathcal{Z}$  so that the
Figure 3: An illustration of the reason why MoVie is effective. We treat the frozen dynamics model as the source of supervision to adapt the latent space of  $h^{\mathrm{SAE}}$  to that of the training views.
Figure 4: The architecture of our spatial adaptive encoder (SAE). We incorporate STNs before the first two convolutional layers of the encoder for better adaptation of the representation  $h$ .
policy would execute the correct mapping  $\mathcal{Z} \mapsto \mathcal{A}$  without training. A vivid illustration is provided in Figure 3.
We freeze the latent dynamics model  $d$ , denoted as  $d^{\star}$ , 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^{\mathrm{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
$$
\mathcal {L} _ {\text {v i e w}} = \left\| d ^ {\star} \left(h ^ {\mathrm {S A E}} (\mathbf {o}), \mathbf {a}\right) - h ^ {\mathrm {S A E}} \left(\mathbf {o} _ {\mathbf {t} + 1}\right) \right\| _ {2}. \tag {2}
$$
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 affine grid and samples features from the original feature map. The point-wise affine transformation is written as
$$
\left( \begin{array}{c} x ^ {s} \\ y ^ {s} \end{array} \right) = \mathcal {T} _ {\phi} (G) = \mathrm {A} _ {\phi} \left( \begin{array}{l} x ^ {t} \\ y ^ {t} \\ 1 \end{array} \right) = \left[ \begin{array}{l l l} \phi_ {1 1} & \phi_ {1 2} & \phi_ {1 3} \\ \phi_ {2 1} & \phi_ {2 2} & \phi_ {2 3} \end{array} \right] \left( \begin{array}{l} x ^ {t} \\ y ^ {t} \\ 1 \end{array} \right) \tag {3}
$$
where  $G$  is the sampling grid,  $(x_i^t,y_i^t)$  are the target coordinates of the regular grid in the output feature map,  $(x_{i}^{s},y_{i}^{s})$  are the source coordinates in the input feature map that define the sample points, and  $\mathrm{A}_{\phi}$  is the affine transformation matrix.
Training strategy for SAE. We use a learning rate  $1 \times 10^{-5}$  for STN layers and  $1 \times 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.
Figure 5: Visualization of three task domains, including DMControl [30], xArm [7], Adroit [20].
Figure 6: Visualization of the training view and four view generalization settings. Trajectory visualization for all tasks is available on our website yangsizhe.github.io/MoVie.
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.
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., DM and 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.
Table 1: Experiment conclusion across all domains and all settings. The best method on each setting is in bold and the relative improvement over TD-MPC is also reported.
Setting	Domain	TD-MPC	DM	IDM+STN	MoVie
Novel view	DMControl	395.61	508.85 (↑29%)	378.80 (↓4%)	623.19 (↑57%)
xArm	16	36 (↑125%)	18 (↑13%)	46 (↑188%)
Adroit	8	14 (↑75%)	12 (↑50%)	34 (↑325%)
Moving view	DMControl	605.98	611.87 (↑1%)	502.34 (↓17%)	673.74 (↑11%)
xArm	20	31 (↑55%)	24 (↑20%)	42 (↑110%)
Adroit	15	20 (↑33%)	15 (↑0%)	45 (↑200%)
Shaking view	DMControl	441.79	348.26 (↓21%)	291.45 (↓26%)	558.23 (↑26%)
xArm	42	44 (↑5%)	44 (↑5%)	45 (↑7%)
Adroit	30	35 (↑17%)	26 (↓13%)	63 (↑110%)
Novel FOV	DMControl	527.47	613.74 (↑16%)	542.43 (↑3%)	770.56 (↑46%)
xArm	34	47 (↑38%)	40 (↑18%)	75 (↑121%)
Adroit	31	51 (↑65%)	31 (↑0%)	68 (↑119%)
All settings	DMControl	492.71	520.68 (↑6%)	426.76 (↓13%)	656.43 (↑33%)
xArm	28	40 (↑43%)	32 (↑14%)	52 (↑86%)
Adroit	21	30 (↑43%)	21 (↑0%)	53 (↑152%)
Table 2: Results in novel view. The best method on each task is in bold.
Novel view	TD-MPC	DM	IDM+STN	MoVie
Cheetah, run	90.13±20.38	254.43±13.89	127.76±16.16	342.39±54.95
Walker, walk	249.34±14.78	262.65±62.86	215.90±29.30	512.71±404.25
Walker, stand	568.01±16.81	635.64±13.70	508.82±15.77	679.90±23.03
Walker, run	127.07±8.06	131.65±2.20	94.49±18.72	94.49±18.87
Cup, catch	922.83±30.74	949.36±9.20	932.75±23.34	961.98±2.68
Finger, spin	137.30±4.47	750.65±9.90	141.03±37.52	892.01±1.20
Finger, turn-easy	380.80±138.54	528.91±160.25	348.00±123.84	705.36±71.98
Finger, turn-hard	261.43±72.83	312.36±134.28	306.76±181.13	331.83±18.98
Pendulum, swingup	40.65±3.72	64.51±2.17	82.46±8.19	528.76±77.51
Reacher, easy	981.25±4.80	982.30±0.83	875.31±150.91	984.66±2.78
Reacher, hard	592.91±59.96	724.91±125.84	533.50±96.90	821.03±67.89
DMControl	395.61	508.85	378.80	623.19 (↑57%)
xArm, reach	3±2	85±13	1±2	95±0
xArm, push	46±5	45±8	60±5	48±7
xArm, peg in box	5±0	5±15	3±2	6±7
xArm, hammer	10±10	10±10	8±5	36±12
xArm	16	36	18	46 (↑188%)
Adroit, door	0±0	10±5	0±0	66±7
Adroit, hammer	6±2	11±10	21±2	11±5
Adroit, pen	18±5	20±8	15±5	25±18
Adroit	8	14	12	34 (↑325%)
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
Table 3: Results in moving view. The best method on each task is in bold.
Moving view	TD-MPC	DM	IDM+STN	MoVie
Cheetah, run	235.37±63.39	344.70±7.54	199.48±25.78	365.22±42.66
Walker, walk	632.77±15.62	707.60±26.72	373.44±30.85	810.19±7.77
Walker, stand	803.97±15.75	706.05±5.75	576.87±44.69	712.48±11.67
Walker, run	295.54±4.39	250.84±7.02	190.02±2.61	281.43±33.34
Cup, catch	887.20±3.67	910.51±4.12	915.16±7.31	951.26±10.68
Finger, spin	636.91±3.91	697.11±8.67	532.90±3.25	896.00±21.65
Finger, turn-easy	715.45±19.23	728.93±180.26	683.85±112.64	744.45±16.54
Finger, turn-hard	593.46±48.22	454.58±88.50	559.20±121.35	558.26±25.64
Pendulum, swingup	26.23±1.85	83.88±10.55	47.20±4.39	236.81±50.23
Reacher, easy	984.10±4.38	981.78±0.60	695.03±236.91	982.58±1.39
Reacher, hard	854.76±31.64	864.55±556.37	752.55±83.55	872.46±51.30
DMControl	605.98	611.87	502.34	673.74 (↑11%)
xArm, reach	15±5	71±2	21±2	80±21
xArm, push	58±5	52±10	55±5	73±5
xArm, peg in box	0±0	0±0	1±2	5±8
xArm, hammer	8±7	0±0	10±5	8±7
xArm	20	31	24	42 (↑110%)
Adroit, door	0±0	10±5	0±0	66±7
Adroit, hammer	25±8	31±12	28±02	43±18
Adroit, pen	20±0	20±10	18±14	26±2
Adroit	15	20	15	45 (↑200%)
Table 4: Results in shaking view. The best method on each task is in bold.
Shaking view	TD-MPC	DM	IDM+STN	MoVie
Cheetah, run	381.04±39.31	317.66±9.57	212.31±19.74	493.54±56.80
Walker, walk	662.13±57.71	627.77±21.47	340.48±36.92	835.99±14.98
Walker, stand	834.70±1.71	604.55±16.20	471.30±69.19	687.96±9.11
Walker, run	251.58±13.69	186.83±3.75	128.31±9.19	291.39±7.94
Cup, catch	752.86±41.81	835.16±32.91	648.75±12.77	951.20±21.01
Finger, spin	89.55±4.72	88.56±5.86	145.68±2.93	284.05±473.73
Finger, turn-easy	717.60±40.93	449.10±20.48	656.23±50.35	694.20±91.02
Finger, turn-hard	411.41±40.28	383.96±32.94	132.40±34.25	629.20±110.57
Pendulum, swingup	23.73±9.24	57.41±4.50	27.30±2.38	48.23±30.25
Reacher, easy	529.58±65.05	225.08±31.02	299.11±59.00	771.38±150.60
Reacher, hard	205.53±25.16	54.80±32.45	144.11±5.60	453.48±381.60
DMControl	441.79	348.26	291.45	558.23 (↑26%)
xArm, reach	76±2	76±7	83±2	88±20
xArm, push	64±3	55±10	61±5	57±13
xArm, peg in box	3±2	15±13	3±2	5±5
xArm, hammer	25±5	28±7	28±12	30±13
xArm	42	44	44	45 (↑7%)
Adroit, door	1±2	10±0	5±5	83±2
Adroit, hammer	58±12	65±18	55±20	70±27
Adroit, pen	30±5	30±8	18±5	35±0
Adroit	30	35	26	63 (↑110%)
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.