PIRLNav: Pretraining with Imitation and RL Finetuning for ObjectNav

In ObjectNav an agent is tasked with searching for an instance of the specified object category (e.g., ‘bed’) in an unseen environment. The agent must perform this task using only egocentric perceptions. Specifically, a RGB camera, Depth sensor²²2We don’t use this sensor as we don’t find it helpful., and a GPS+Compass sensor that provides location and orientation relative to the start position of the episode. The action space is discrete and consists of move_forward ($0.25m$), turn_left ($30^{\circ}$), turn_right ($30^{\circ}$), look_up ($30^{\circ}$), look_down ($30^{\circ}$), and stop actions. An episode is considered successful if the agent stops within $1m$ Euclidean distance of the goal object within $500$ steps and is able to view the object by taking turn actions [14].

We use scenes from the HM3D-Semantics v0.1 dataset [16]. The dataset consists of $120$ scenes and $6$ unique goal object categories. We evaluate our agent using the train/val/test splits from the 2022 Habitat Challenge³³3https://aihabitat.org/challenge/2022/.

Ramrakhya et al. [1] collected ObjectNav demonstrations for the Matterport3D dataset [15]. We begin our study by replicating this effort and collect demonstrations for the HM3D-Semantics v0.1 dataset [16]. We use Ramrakhya et al.’s Habitat-WebGL infrastructure to collect $77k$ demonstrations, amounting to ${\sim}2378$ human annotation hours.

We use behavior cloning to pretrain our ObjectNav policy on the human demonstrations we collect. Let $\pi^{BC}_{\theta}(a_{t}\mid o_{t})$ denote a policy parametrized by $\theta$ that maps observations $o_{t}$ to a distribution over actions $a_{t}$. Let $\tau$ denote a trajectory consisting of state, observation, action tuples: $\tau=\big{(}s_{0},o_{0},a_{0},\ldots,s_{T},o_{T},a_{T}\big{)}$ and $\mathcal{T}=\big{\{}\tau^{(i)}\big{\}}_{i=1}^{N}$ denote a dataset of human demonstrations. The optimal parameters are

We use inflection weighting [43] to adjust the loss function to upweight timesteps where actions change (i.e. $a_{t-1}\neq a_{t}$).

Our ObjectNav policy architecture is a simple CNN+RNN model from [30]. To encode RGB input $(i_{t}=$ CNN$(I_{t}))$, we use a ResNet50 [44]. Following [30], the CNN is first pre-trained on the Omnidata starter dataset [45] using the self-supervised pretraining method DINO [46] and then finetuned during ObjectNav training. The GPS+Compass inputs, $P_{t}=(\Delta x,\Delta y,\Delta z)$, and $R_{t}=(\Delta\theta)$, are passed through fully-connected layers $p_{t}=$ FC$(P_{t}),r_{t}=$ FC$(R_{t})$ to embed them to 32-d vectors. Finally, we convert the object goal category to one-hot and pass it through a fully-connected layer $g_{t}=$ FC$(G_{t})$, resulting in a 32-d vector. All of these input features are concatenated to form an observation embedding, and fed into a 2-layer, 2048-d GRU at every timestep to predict a distribution over actions $a_{t}$ - formally, given current observations $o_{t}=[i_{t},p_{t},r_{t},g_{t}]$, $(h_{t},a_{t})=$ GRU$(o_{t},h_{t-1})$. To reduce overfitting, we apply color-jitter and random shifts [47] to the RGB inputs.

The RL objective is to find a policy $\pi_{\theta}(a|s)$ that maximizes expected sum of discounted future rewards. Let $\tau$ be a sequence of object, action, reward tuples ($o_{t}$, $a_{t}$, $r_{t}$) where $a_{t}\sim\pi_{\theta}(\cdot\mid o_{t})$ is the action sampled from the agent’s policy, and $r_{t}$ is the reward. For a discount factor $\gamma$, the optimal policy is

To solve this maximization problem, actor-critic RL methods learn a state-value function $V(s)$ (also called a critic) in addition to the policy (also called an actor). The critic $V(s_{t})$ represents the expected value of returns $R_{t}$ when starting from state $s_{t}$ and acting under the policy $\pi$, where returns are defined as $R_{t}=\sum_{i=t}^{T}\gamma^{i-t}r_{i}$. We use DD-PPO [48], a distributed implementation of PPO [42], an on-policy RL algorithm. Given a $\theta$-parameterized policy $\pi_{\theta}$ and a set of rollouts, PPO updates the policy as follows. Let $\hat{A_{t}}=R_{t}-V(s_{t})$, be the advantage estimate and $p_{t}(\theta)=\frac{\pi_{\theta}(a_{t}|o_{t})}{\pi_{\theta_{\text{old}}}(a_{t}|o_{t})}$ be the ratio of the probability of action $a_{t}$ under current policy and under the policy used to collect rollouts. The parameters are updated by maximizing:

We use a sparse success reward. Sparse success is simple (does not require hyperparameter optimization) and has fewer unintended consequences (e.g. Maksymets et al. [17] showed that typical dense rewards used in ObjectNav actually penalize exploration, even though exploration is necessary for ObjectNav in new environments). Sparse rewards are desirable but typically difficult to use with RL (when initializing training from scratch) because they result in nearly all trajectories achieving $0$ reward, making it difficult to learn. However, since we pretrain with BC, we do not observe any such pathologies.

We use the behavior cloned policy $\pi^{BC}_{\theta}$ weights to initialize the actor parameters. However, notice that during behavior cloning we do not learn a critic nor is it easy to do so – a critic learned on human demonstrations (during behavior cloning) would be overly optimistic since all it sees are successes. Thus, we must learn the critic from scratch during RL. Naively finetuning the actor with a randomly-initialized critic leads to a rapid drop in performance⁴⁴4After the initial drop, the performance increases but the improvements on success are small. (see Fig. 8) since the critic provides poor value estimates which influence the actor’s gradient updates (see Eq.(3)). We address this issue by using a two-phase training regime:

Phase 1: Critic Learning. In the first phase, we rollout trajectories using the frozen policy, pre-trained using BC, and use them to learn a critic. To ensure consistency of rollouts collected for critic learning with RL training, we sample actions (as opposed to using argmax actions) from the pre-trained BC policy: $a_{t}{\sim}\pi_{\theta}(s_{t})$. We train the critic until its loss plateaus. In our experiments, we found $8M$ steps to be sufficient. In addition, we also initialize the weights of the critic’s final linear layer close to zero to stabilize training.

Phase 2: Interactive Learning. In the second phase, we unfreeze the actor RNN⁵⁵5The CNN and non-visual observation embedding layers remain frozen. We find this to be more stable. and finetune both actor and critic weights. We find that naively switching from phase 1 to phase 2 leads to small improvements in policy performance at convergence. We gradually decay the critic learning rate from $2.5\times 10^{-4}$ to $1.5\times 10^{-5}$ while warming-up the policy learning rate from $0$ to $1.5\times 10^{-5}$ between $8M$ to $12M$ steps, and then keeping both at $1.5\times 10^{-5}$ through the course of training. See Fig. 3. We find that using this learning rate schedule helps improve policy performance. For parameters that are shared between the actor and critic (i.e. the RNN), we use the lower of the two learning rates (i.e. always the actor’s in our schedule). To summarize our finetuning methodology:

1. –

   First, we initialize the weights of the policy network with the IL-pretrained policy and initialize critic weights close to zero. We freeze the actor and shared weights. The only learnable parameters are in the critic.

2. –

   Next, we learn the critic weights on rollouts collected from the pretrained, frozen policy.

3. –

   After training the critic, we warmup the policy learning rate and decay the critic learning rate.

4. –

   Once both critic and policy learning rate reach a fixed learning rate, we train the policy to convergence.

Comparing with the RL-finetuning approach in VPT [21]. We start by comparing our proposed RL-finetuning approach with the approach used in VPT [21]. Specifically, [21] proposed initializing the critic weights to zero, replacing entropy term with a KL-divergence loss between the frozen IL policy and the RL policy, and decay the KL divergence loss coefficient, $\rho$, by a fixed factor after every iteration. Notice that this prevents the actor from drifting too far too quickly from the IL policy, but does not solve uninitialized critic problem. To ensure fair comparison, we implement this method within our DD-PPO framework to ensure that any performance difference is due to the fine-tuning algorithm and not tangential implementation differences. Complete training details are in the Sec. C.3. We keep hyperparameters constant for our approach for all experiments. Table 1 reports results on HM3D val for the two approaches using $20k$ human demonstrations. We find that PIRLNav achieves $+2.2\%$ Success compared to VPT and comparable SPL.

Ablations. Next, we conduct ablation experiments to quantify the importance of each phase in our RL-finetuning approach. Table 2 reports results on the HM3D val split for a policy BC-pretrained on $20k$ human demonstrations and RL-finetuned for $300M$ steps, complete training details are in Sec. C.4. First, without a gradual learning transition (row $2$), i.e. without a critic learning and LR decay phase, the policy improves by $1.6\%$ on success and $8.0\%$ on SPL. Next, with only a critic learning phase (row $3$), the policy improves by $4.7\%$ on success and $7.1\%$ on SPL. Using an LR decay schedule only for the critic after the critic learning phase improves success by $7.4\%$ and SPL by $6.3\%$, and using an LR warmup schedule for the actor (but no critic LR decay) after the critic learning phase improves success by $6.2\%$ and SPL by $6.1\%$. Finally, combining everything (critic-only learning, critic LR decay, actor LR warmup), our policy improves by $9.9\%$ on success and $7.3\%$ on SPL.

ObjectNav Challenge 2022 Results. Using our overall two-stage training approach of BC-pretraining followed by RL-finetuning, we achieve state-of-the-art results on ObjectNav– $65.0\%$ success and $33.0\%$ SPL on both the test-standard and test-challenge splits and $70.4\%$ success and $34.1\%$ SPL on val. Table 3 compares our results with the top-4 entries to the Habitat ObjectNav Challenge 2022 [50]. Our approach outperforms Stretch [24] on success rate on both test-standard and test-challenge and is comparable on SPL ($1\%$ worse on test-standard, $4\%$ better on test-challenge). ProcTHOR [49], which uses $10k$ procedurally-generated environments for training, achieves $54\%$ success and $32\%$ SPL on test-standard split, which is $11\%$ worse at success and $1\%$ worse at SPL than ours. For sake of completeness, we also report results of two unpublished entries uploaded to the leaderboard – Populus A. and ByteBOT. Unfortunately, there is no associated report yet with these entries, so we are unable to comment on the details of these approaches, or even whether the comparison is meaningful.

Using the BC setup described in Sec. 3.3, we train on SP, FE, and HD demonstrations. Since these demonstrations vary in trajectory length (e.g. SP are significantly shorter than FE), we collect ${\sim}12M$ steps of experience with each method. That amounts to $240k$ SP, $70k$ FE, and $77k$ HD demonstrations respectively. As shown in Table 4, BC on $240k$ SP demonstrations leads to $6.4\%$ success and $5.0\%$ SPL. We believe this poor performance is due to an imitation gap [51], i.e. the shortest path demonstrations are generated with access to privileged information (ground-truth map of the environment) which is not available to the policy during training. Without a map, following the shortest path in a new environment to find a goal object is not possible. BC on $70k$ FE demonstrations achieves $44.9\%$ success and $21.5\%$ SPL, which is significantly better than BC on shortest paths ($+38.5\%$ success, $+16.5\%$ SPL). Finally, BC on $77k$ HD obtains the best results – $64.1\%$ success, $27.1\%$ SPL. These trends suggest that task-specific exploration (captured in human demonstrations) leads to much better generalization than task-agnostic exploration (FE) or shortest paths (SP).

Using the BC-pretrained policies on SP, FE, and HD demonstrations as initialization, we RL-finetune each using our approach described in Sec. 4. These results are summarized in Fig. 4. Perhaps intuitively, the trends after RL-finetuning follow the same ordering as BC-pretraining, i.e. RL-finetuning from BC on HD $>$ FE $>$ SP. But there are two factors that could be leading to this ordering after RL-finetuning – 1) inconsistency in performance at initialization (i.e. BC on HD is already better than BC on FE), and 2) amenability of each of these initializations to RL-finetuning (i.e. is RL-finetuning from HD init better than FE init?).

We are interested in answering (2), and so we control for (1) by selecting BC-pretrained policy weights across SP, FE, and HD that have equal performance on a subset of train $={\sim}48.0\%$ success. This essentially amounts to selecting BC-pretraining checkpoints for FE and HD from earlier in training as ${\sim}48.0\%$ success is the maximum for SP.

Fig. 5 shows the results after BC and RL-finetuning on a subset of the HM3D train and on HM3D val. First, note that at BC-pretraining train success rates are equal ($={\sim}48.0\%$), while on val FE is slightly better than HD followed by SP. We find that after RL-finetuning, the policy trained on HD still leads to higher val success ($66.1\%$) compared to FE ($51.3\%$) and SP ($43.6\%$). Notice that RL-finetuning from SP leads to high train success, but low val success, indicating significant overfitting. FE has smaller train-val gap after RL-finetuning but both are worse than HD, indicating underfitting. These results show that learning to imitate human demonstrations equips the agent with navigation strategies that enable better RL-finetuning generalization compared to imitating other kinds of demonstrations, even when controlled for the same BC-pretraining accuracy.

Results on SP-favoring and FE-favoring episodes. To further emphasize that imitating human demonstrations is key to good generalization, we created two subsplits from the HM3D val split that are adversarial to HD performance – SP-favoring and FE-favoring. The SP-favoring val split consists of episodes where BC on SP achieved a higher performance compared to BC on HD, i.e. we select episodes where BC on SP succeeded but BC on HD did not or both BC on SP and BC on HD failed. Similarly, we also create an FE-favoring val split using the same sampling strategy biased towards BC on FE. Next, we report the performance of RL-finetuned from BC on SP, FE, and HD on these two evaluation splits in Table 5. On both SP-favoring and FE-favoring, BC on HD is at $0\%$ success (by design), but after RL-finetuning, is able to significantly outperform RL-finetuning from the respective BC on SP and FE policies.

In this section, we investigate how BC-pretraining $\rightarrow$ RL-finetuning success scales with no. of BC demonstrations.

Human demonstrations. We create HD subsplits ranging in size from $2k$ to $77k$ episodes, and BC-pretrain policies with the same set of hyperparameters on each split. Then, for each, we RL-finetune from the best-performing checkpoint. The resulting BC and RL success on HM3D val vs. no. of HD episodes is plotted in Fig. 1. Similar to [1], we see promising scaling behavior with more BC demonstrations.

Interestingly, as we increase the size of of the BC pretraining dataset and get to high BC accuracies, the improvements from RL-finetuning decrease. E.g. at $20k$ BC demonstrations, the BC$\rightarrow$RL improvement is $10.1\%$ success, while at $77k$ BC demonstrations, the improvement is $6.3\%$. Furthermore, with $35k$ BC-pretraining demonstrations, the RL-finetuned success is only $4\%$ worse than RL-finetuning from $77k$ BC demonstrations ($66.4\%$ vs. $70.4\%$). Both suggest that by effectively leveraging the trade-off between the size of the BC-pretraining dataset vs. performance gains after RL-finetuning, it may be possible to achieve close to state-of-the-art results without large investments in demonstrations.

How well does FE Scale? In Section 5.1, we showed that BC on human demonstrations outperforms BC on both shortest paths and frontier exploration demonstrations, when controlled for the same amount of training experience. In contrast to human demonstrations however, collecting shortest paths and frontier exploration demonstrations is cheaper, which makes scaling these demonstration datasets easier. Since BC performance on shortest paths is significantly worse even with $3$x more demonstrations compared to FE and HD ($240k$ SP vs. $70k$ FE and $77k$ HD demos, Sec. 5.1), we focus on scaling FE demonstrations. Fig. 6 plots performance on HM3D val against FE dataset size and a curve fitted using $75k$ demonstrations to predict performance on FE dataset-sizes $\geq$ $75k$. We created splits ranging in size from $10k$ to $150k$. Increasing the dataset size doesn’t consistently improve performance and saturates after $70k$ demonstrations, suggesting that generating more FE demonstrations is unlikely to help. We hypothesize that the saturation is because these demonstrations don’t capture task-specific exploration.