Accelerating Visual Sparse-Reward Learning with Latent Nearest-Demonstration-Guided Explorations
Authors: Ruihan Zhao
ufuk topcu
Sandeep P. Chinchali
Mariano Phielipp
Venue: CoRL 2024 (2024)
Keywords: Computer Vision
Sparse Reward
RL from Demonstrations
TL;DR: Under-an-hour RL training of real robot manipulation is possible with demonstration-guided reward augmentation.

Abstract: Recent progress in deep reinforcement learning (RL) and computer vision enables artificial agents to solve complex tasks, including locomotion, manipulation, and video games from high-dimensional pixel observations. However, RL usually relies on domain-specific reward functions for sufficient learning signals, requiring expert knowledge. While vision-based agents could learn skills from only sparse rewards, exploration challenges arise. We present Latent Nearest-demonstration-guided Exploration (LaNE), a novel and efficient method to solve sparse-reward robot manipulation tasks from image observations and a few demonstrations. First, LaNE builds on the pre-trained DINOv2 feature extractor to learn an embedding space for forward prediction. Next, it rewards the agent for exploring near the demos, quantified by quadratic control costs in the embedding space. Finally, LaNE optimizes the policy for the augmented rewards with RL. Experiments demonstrate that our method achieves state-of-the-art sample efficiency in Robosuite simulation and enables under-an-hour RL training from scratch on a Franka Panda robot, using only a few demonstrations.

Accelerating Visual Sparse-Reward Learning with
Latent Naares-DDemonsraaiinn-nniied Explorations
Ruihan Zhao!
Ufuk Topcu't
Sandeep Chinchali
Mariano Phielipp
'The University of Texas at Austin
Intel Al Lab
Recent progress in deep reinforcement learning (RL) and computer vision enables artificial agents to solve complex tasks, including locomotion, ma- nipulation, and video games from high-dimensional pixel observations. How- ever, RL usually relies on domain-speciff reward functions for sufficient learn- ing signals, requiring expert knowledge. While vision-based agents could learn skills from only sparse rewards, exploration challenges arise. We present Latent priress-demonstration -ruiidd Exploration (LaNE), novel and efficient method solve sparse-reward robot manipulation tasks from image observations and few demonstrations. First, LaNE builds the pre-trained DINOv2 feature extrac- tor to learn embedding space for forward prediction. Next, it rewards the agent for exploring near the demos, quantified by quadratic control costs in the embed- ding space. Finally, LaNE optimizes the policy for the augmented rewards with RL. Experiments demonstrate that our method achieves state-of-the-an sample efficiency in Robosuite simulation and enables under-an-hour RL training from scratch on Franka Panda robot, using only a few demonstrations.
Abstract:
Keywords: Computer Vision, Sparse Reward, RI from Demonstrations
Reach Location 1 Demo 21 minutes
Block 5 Demos
Open Drawer 1 Demo
Insert Demos minutes
45 minutes
52 minutes
Figure 1: LaNE enables under-an-howe RI training from scratch on a Franka Panda arm from image observa- tions and sparse rewards, utilizing only few demonstrations. LaNE achieves unparalleled sample efficiency hy learning an embedding space to quantify state proximity and reward explorations close to the demonstrations.
1 Introduction
Deep reinforcement learning (RL) is a versatile approach that learns from interaction data, often without an explicit, hand-coded dynamics model. Through environmental interactions, RL agents can learn optimal policies from dense or sparse reward feedback. State-of-the-ard approaches al- low learning policies for discrete actions and continuous action spaces while taking either low- dimensional state vectors or high-dimensional sensor readings [1, 3,
However, applying deep RL to real-life domains, including real-hardware robot learning, remains difficult. One challenge is the need to reliably track the complete system state [5]. A policy that directly maps images to optimal actions could alleviate such engineering challenges: data augmen- tation and self-supervised learning have enabled policy learning from image observations with high
Equal advising. Prooeect page:
8th Conference Robot Learning (CoRL. 2024), Munich, Germany.

---

sample efficiency [6. 7, 8, 9]. Meanwhile, it is also hard to assign informative rewards in scalable way. Reward engineering requires domain-specific knowledge: popular simulated environments provide optional dense reward functions based on heuristics [10, 11] but rely on state readings not readily available in the real physical system. Thus, there has been a lot of effort to help RL agents explore more effectively in environments with sparse or no rewards [12, 13, 14].
We present Latent Yearest-demonstratioe guided Exploration (LaNE), novel approach to tackle the exploration challenge in image-baser control tasks with sparse rewards. LaNE builds on the idea of learning from demonstrations (I.fD) and draws inspiration from reward shaping (15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25J. Our main insight that each step in the demonstration can be considered a subgoal, and the agent should be credited for reaching a similar state (see Fig. 2). Our approach presents two main contributions:
1. We define distance measure among image observations by learning a lower-dimensional latent space. Specifically, we train a Variational Autoencoder (VAE) so that the resulting latent space's forward dynamics are locally linear. A quadratic control cost in this space effectively identifies nearby states, whereas pre-trained embeddings from even a state-of- the-art computer vision model are insufficient.
2.
We propose systematic way to provide dense reward signals in sparse-reward tasks un- der the LID paradigm. When exploring near the demonstrations, the RL agent receives additional task-progress--nformedd rewards. The augmented reward function also derives bounded value functions, significantly improving training stability.
LaNE is independent of the underlying RL algorithm, and use Soft Actor- Critic [1] in this paper. In Robosuite simulation [11], our method significantly improves sample efficiency when learning morizon, sparse-reward visual manipulation tasks. On real hardware, LaNE enables learning various manipulation tasks from scratch with a Franka Panda arm, each with under an hour of train- ing and only one to five demonstrations (see Fig. 1).
2 Related Work
Exploration known challenge in deep RL, especially in sparse-reward environments [26]. One solution is to guide exploration with domain knowledge. In reward shaping, intermediate rewards can be added at important checkpoint locations [11], derived from physics knowledge [27]. or learned from human annotations [28]. Aside from external guidance, other works aim to improve the agents' intrinsic exploration behavior. Maximum entropy RL balances exploration and exploitation by encouraging high policy entropy [1]. Hierarchical RI. and intrinsically motivated skill learning are also potential remedies [12, 29]. Our work aims to alleviate the need for domain expertise in the reward-shaping approach by using few demonstrations.
Learning from demonstrations (LfD) has proven helpful in expediting RL, especially in sparse ward settings. Prior works have introduced various auxiliary training objectives. Behavior cloning
[21, 23] and supervised Q-value updates [18] can effectively learn from optimal demonstrations.
[30] employ an information theoretic approach to guide policy distributions. GAIL [31], AIRL [32] and DAC [33] perform adversarial training to distinguish expert and policy rollouts. [17] and [20] learn potential functions from value estimates, allowing imperfect or mismatched demonstrations to be used. CoDER [34] performs contrastive learning to pre-train the image encoders. Other methods
perform warm start with scripted behavior-cloned policies [14, 22]. Our method adopts the LiD paradigm for efficient learning and is compatible with sub-optimal demonstrations.
Aside from more efficient exploration and LiD, model-based RL methods aim to improve efficiency by generating new synthetic experiences with learned world models. DreamerV2 (6) learns an ac- curate discrete world model from high-dimensional image inputs, enabling human-level RL per- formance in Atari games. Modem [35] combines model-based learning with demonstrations by over sampligm demonstrated data to form behavior prior. Our method LaNE also learns forward

---

%
of
Demo
Policy
%
ob
oi
(o, s) a(e)
Slice
Nearest
Replay Buffer
New Experience
HIOL
(e) Demonstrations
(b) Nearest-Dema Metching
(c) RL Updetes
Figure 2: Latent dearest-demonstaationn -uuiidd Exploration (LaNE) augments the sparse task reward with dense exploration reward in vision-based RL. (Left) We utilize variable-length demonstrations, each consisting of observations of and actions aj. (Middle) A dense exploration reward is given when transition lands sufficiently close demonstration and is discounted based on its distance to the goal. (Right) Using the combined reward signal, the RL agent learns to map sensor observation
action
dynamics model but only as an auxiliary task for representation learning. Hence, it can be integrated with both model-based and model-free RL methods.
Finally, a rich collection of prior work has studied ways to learn better representations from high- dimensional image observations. RAD and DrQ [8, 9, 36] perform data augmentation on image observations to promote task-relevant features. CURL and CoDER [7, 34] use contrastive learning as a self-super auxiliary objective for feature learning. TCC and TCN [37, 38] perform self- supervised learning from video sequences. Recently, foundation models have also proven beneficial for learning generalizable agents [39]. Our method builds on state-of-the-arr foundation model, DINOv2 [40], and uses data augmentation to learn robust embedding space.
3 Problem Setting
We tackle the challenges of vision-based RI. in long-horizon, sparse reward tasks, where the agent only receives positive reward at task completion while getting a constant negative reward Thive everywhere else. The interpretation of such a reward function is that the agent gets a high reward only at task completion but is penalized for the trajectory length. Formally, letting & denote the set of goal states, we define the reward function follows:
r(s, a, 8')
Idoue
0 S'E
Flive
otherwise.
The sparse reward function reduces the need for expert knowledge or human intervention, making it much easier to implement in real-world environment, but makes exploration hard in training.
Method
We present Latent dearest-demonstrationn -uuiddd Exploration, an efficient RI. algorithm centered around a few demonstrations to tackle the exploration challenge with sparse-reward learning. The core idea is provide additional task-progress-aware dense rewards when the agent close the demonstrations. We learn structured embedding space to quantify state proximity from image observations by learning a latent dynamics model as an auxiliary task. The cleverly designed aug- mented reward function also derives bounded value functions, enabling us to perform value clipping and greatly enhance training stability.
4.1 Reinforcement Learning from Demonstrations
LaNE utilizes demonstration set D consisting of successful trajectories of observations and actions. Each demonstration trajectory may have different length T,, but must terminate in the goal set G. We assume the demonstrations to come from human operator or heuristic controller

---

and, thus, can be sub sub-optimal. We formalize the notations as follows: T denotes the trajectory for demonstration i. Sr is the underlying true environment state, and Og is the high-dimensional observation, such as images. ag is the action taken by the demonstrator. Note that the RL algorithm cannot access the ground truth state 8, in the demonstrations, but only observations, as shown below:
D = [T|, .T....T) }
T = (ob: ap: of ... -
(2)
&T EG.
The demonstrations are stored in two forms. First, the trajectory form records the steps to success from each state, allowing us to discount the exploration reward according to task progress, as detailed in Section 4.3. Next, they are sliced into experience tuples (o, a,,o' r; d) and placed in replay buffer for representation learning and off- ofl-pol RL updates. Here, o' denotes the next observation after o. r is the sparse reward as defined in Equation and d is a Boolean variable indicating episode termination. LaNE uses iottinn-irsttouu replay buffer B with capacity for limited number of transitions, but the demonstrations always retained ensure sufficient reward signal.
4.2
Dongmentation-Invaiin Distance Measure
The key technical challenge behind LaNE is to find a demonstration state closest to the agent's current state and quantify proximity. Com- puting the distance between two states from their respective image observations is nontriv- ial: two drastically different states might only differ by a few pixels. Conversely, the same underlying state might appear very different in two images due to task-irrelevant background features. To solve this, we embed the images into a low-dimensional latent space obtain viable distance measure. Inspired by Embed to Control (E2C) [41], we train a VAE [42] and enforee locally linear dynamics model to regularize the structure of the latent space. The locally-linear dynamics model captures our goal for the latent space to be temporally con- sistent since we have a multi-step control task.
Frozen
Frozen
ZAONIa
ZAONTa
a
w
Predict
E
Dg
A B
Ma
2
c
Figure 3: LaNE learns latent space with locally linear dynamics. Given transition tuple (o, o'), the ob- servations
first encoded by a frozen DINOv2 model into w and w'. Next, the encoder Egg further embeds w into a oiw-dimensional latent state 2. The forward model Mş. predicts the transition matrices A, B and offset Finally, the decoder Do reconstructs w² from the predicted a', where I' A + Ba + c. The train- able modules Eg, Mg. and Do are colored in orange.
Our method differs from E2C in three key ways: 1. LaNE leverages the feature extractor from strong pre-trained image model by embedding and predicting the frozen DINOv2 [40] features. As we find with ablation studies in Section 5.2, learning on top of DINOv2 features w is superior to learning directly from pixels o. 2. We learn a latent space robust to pixel-space perturbations. has been shown that data augmentation is crucial for efficient and robust image- based reinforcement learning [8, 9, 34, 36]. Hence, we apply a random dagm-auggennation function f(-) during represen- tation learning. 3. Our RL policy network
separate CNN encoder to learn in an unconstrained embedding space and benefit from low inference latency. Overall, the trainable components include the encoder Eg. the decoder Dg. and the transition model Mg.
LaNE optimizes a variational lower bound (ELBO) objective across transition tuples (o, a, o'); sam- pled from the replay buffer. We assume the latent states z form unit Gaussian prior p(z) \O Z). The encoded distributions q(z /w) and decoded distribution p(w 22) are also modeled with Gaussian distributions. During training, the data-augmented observation features are encoded into their latent distributions whose mean μ and diagonal covariance matrix Σ are predicted by the encoder network Es:
(z(z) w) = (N,, E), where (n, E) = Es (w), w DINOv2(((o)) z' ~ (6 |w') = -N E'), where (μ', E') = Es(w),, w' DINOv2((((')).
(3)

---

The one-step forward model in the latent space is locally linear in the state and action, whose param- eters (matrices A, B and offset c) depend on the starting state, as predicted by the latent transition model Mg. Prior work shows that a latent linear dynamics model is tractable to learn but provides modeling flexibility through local linearity [41]. The linear transition model allows the prediction of the next step latent distribution using the current distribution and action as follows:
z' Az + Ba + C, where (A, B) M²(z)
(4)
(w 2, () (n²,, E), where ji' = Au + Ba + c, AEA
(5)
Finally, the decoder Do reconstructs the next step observation embedding from the predicted next step latent vector:
w = Do(2").
(6)
The encoder Es, decoder Da and transition model Ma are updated jointly using a combined loss with three terms. First, we want the sampled starting latent state z to be reconstructed back to the original image features w. Similarly, as we pass the sample through the dynamies model, the resulting latent state prediction 2' should be reconstructed back to w'. Finally, to ensure the latent dynamics model is consistent over multiple steps, we want the predicted distribution go (2//w, a) and encoded distribution (z (z'lw") to be similar. Formally, we write the overall training objective L as follows, where À is a hyper-parameter for weighing the two loss terms:
CELOO
E go- z'ago
log p(w])2) log pepp((w12"
+ DKL (as(z) w)
p(z)
(7)
Ldynamics
= E zruge
DKL qu(zz) 2.
90(z'| (iallw')
(8)
C
=
E (o,0,d))EB
LELBO + XLdynamiic
(9)
In essence, we minimize the reconstruction error for the VAE using the ELBO objective (term 1) and the forward prediction error in the latent space using KI. divergence (term 2).
The learned latent space allows us to define a dynamics-awart distance measure between observa-
tions. Specifically, for two observations on and 02, we define the Augmentation-invarian Distance Mcasure (ADM) to be the root quadratic cost between the augmented and encoded states 21 and 22:
d(o11 02) == ((21 2² (a 22))²
(10)
Our design echoes the quadratic cost function commonly used in optimal control. Using an identity weighting matrix Q further simplifies ADM to the Euclidean distance in the latent space. We apply this simplification in our experiments, but other weighting matrices could be useful when the agent observes both images and proprioceptive states.
4.3 Demonstration-CGiided Exploration
We propose a systematic eeward-engineering approach to credit the agent for staying close to demon- strations. Given an experience tuple (o, a, o', r, d),, we assign an additional exploration reward r of is sufficiently close demonstrated state, up to a distance threshold e, which is dynamically computed. We define as the average distance between consecutive demonstration observations:
== E. [d(ot, o(+1)],
oi. E D.
(11)
The threshold approximates the distance of one environment step. As the agent gathers new periences, we re-compute This is necessary because the encoder Es, decoder Dø, and dynamics model parameters My are constantly updated with the latest experience to ensure that ADM is not overfitted to only the demonstration data. In addition, we find the trajectory index and time step t of the nearest demonstration using the ADM d:
/ arg min d(o' ,o)), of D.
(12)
1,t
5

---

Lift Block
Open Door 10 demes
Stack Blocks
Move Can 20 demos
5 demos
10 demos
Figure 4: LaNE achieves state-of-the-art sample efficiency in four Robosuite visual manipulation tasks. The RI. agent observes RGB images from cameras, one front (shown above) and other on the gripper.
After we compute c, ;", and t", we assign a dense reward Idense by augmenting the environment reward with dask-progress-infommen exploration reward:
a Ti--tt re
of(dd E
ddnase
(13)
otherwise.
Here, the exploration reward rg T.- is modulated by the expected step to success, which the difference between the demo trajectory length TІ- and t*. In the context of LfD, r can be interpreted as point estimate of the potential function o', echoing prior work in this domain [17, 20]. The discounting factor a is sypee--arammttr chosen independently from the RL discounting factor y, and the nominal exploration reward Te is constant.
Inspecting the augmented reward Idense. we see that when of finds its nearest neighbor close to one of the successful terminal states, o' is awarded almost the full nominal exploration reward re- When of is close to one the earlier steps in a demonstration, Te is heavily discounted. Finally, if we are very far from any demonstration observation (relative to the distance threshold e), or we are the goal, the RL agent gets only the environment reward (case 2 in Eq. 13). LaNE is versatile because can train control policy to maximize Tdense using any off-the-shelf RL algorithm.
4.4 Prioritized Replay and Value Clipping
We improve training efficiency and stability by performing prioritized replay and Q-value clipping- Prioritized replay is a standard tool when learning from demonstrations [19]: each batch transitions, we sample least pa fraction from the demonstrations, where and pa are hyper- parameters. Conservative q-value estimates are also widely used stabilize training [1, 43]. LaNE stands out because Tdense derives upper and lower bounds the q-value landscape, allowing to use clipped value target when performing temporal difference updates.
The definition of Fdense in Eq. 13 contains the nominal exploration reward re. constant hyper- parameter. We pick [rel < |Nivvl to obtain bounded Q-values. Under Idenses a transition either receives positive reward Tdone and terminates the episode or receives a non-iisitive reward Flive + 1. r The highest Q-value achieved at task completion. On the other hand, in the worst where the agent receives Flive all the time and never succeeds, the Q-value is bounded below by Er Y' Nive -ivvee Thus, we obtain the following bound:
1-9 Flive S 2((o, a) S Tdone-
5 Experiments
5.1 Simulated Manipulation Tasks
We solve four robot manipulation tasks from the Robosuite simulator [11]: lifting block, opening door, stacking blocks, and moving soda can. The RL agent observes 128 128 RGB images from two cameras, mounted on the robot gripper and one in front of the robot. The robot uses Operational Space Control the agent predicts actions to change the robot's hand displacement,

---

EXSMPCC
NoDon
CaDZA
LINE
L/f Block
Open a Doar
Stack Steds
More Can
1.0
0.8
Bake
0.6
0.6
2.8
5000000
02
0.0
0.0
0.8
0.5
05
anstrnmment stapa
led
145
Figure 5: We compare LaNE with optimal control and state-of-the-art RL methods: E2C [41], MoDem [35] and CoDER [34], Our methad (red) consistently learns faster and converges to higher success rates than all three baseline methods.
RAD
N6.
RAD+NT
RAD+FB-WK
LIVE
Ne 3MM0
DIMO Oiyy
LikE
Stack Blocks
Stack Blocks
Moue
1.0
1.0
2.0
1.0
520000 Rate
0.8
Su000ss Nate
Bate
0.0
0.00
2.6
2.5
0.6
0.4
-4
0.1
0.2
222
02
0.0
0.0
0.0
0.0
0.0
0.5
2.0
1.5
0.0
5.5
10
15 2.0
2.5
0.0
0.5
0.0
0.5
1.3
2.5
2.0
2.5
steps
steps
Figure 6: Left: Starting from RAD [8], we add PR (prioritized replay). VC (value clipping), and exploration reward re to understand their contributions. PR and VC necessary for stable learning but inefficient. re by itself results unstable training. Hence, all components of LaNE are necessary. Right: We experiment with LaNE variants without pre- vre-training (No DINO) or fine-tuning (DINO Only). While either variant improves from the baseline where no exploration reward is given, combining both results in the quickest exploration.
rotation, and gripper width for 7 degrees of freedom. Figure shows one of the camera angles. The demonstrations come from hand-coded controller that utilizes state information. The state is unavailable for the RI. agent, who must learn from images alone.
We compare LaNE with optimal control, model-free, and model-based RL. For E2C [41], we run an MPC controller in the latent space to minimize a quadratic state cost towards the demonstrated goal state. CoDER [34], and MoDem [35] are tfate-of-the-art vision-based RL algorithms in the LiD setting. We initiate all methods with the same demonstrations and measure the evaluation success rates during training. Figure 5 shows the mean and standard deviation across five random seeds as function of training environment steps. Our method outperforms all three baselines across all tasks while showing major advantages in the two more challenging tasks.
5.2 Ablation Studies
Our approach utilizes multiple techniques to maximize the utility of demonstrations and to speed up learning. We perform ablation studies to answer two key questions: 1. Is reward augmentation really necessary, or are regularization tricks like prioritized replay (PR) and value clipping (VC) by themselves sufficient? 2. LaNE learns embedding space by fine-tuning pre-trained computer vision model. Are both the fine-tuning step and the pre-trained DINOv2 model necessary?
For question 1, we start from a standard image-based RL algorithm RAD [8] (initialized with demon- strations) and add our key components, namely prioritized replay (PR), value clipping (VC), and the exploration reward dedese As shown by the left half of Figure 6, importance sampling and value clipping help stabilize training but are not efficient enough. Ideuse significantly contributes on top and allows the robot complete the long-horizon task reliably much earlier in training.
For question 2, we experiment with two variants of LaNE, one without DINOv2 and one without fine-tuning. When LaNE trains without DINOv2, it initializes the encoder and decoder from scratch and learns to predict the images directly. When LaNE uses the pre-trained model and skips the fine- tuning steps, we directly use the Fuclidean distance between their respective DINOv2 embedding vectors. The results in the right side of Figure demonstrate that performing reward augmenta- tion without pre-trained DINOv2 fine-tuning helps, but combining both allows the most efficient learning. Notably, DINOy2 improves training convergence, as showcased in Figure 9.
7

---

Figure 7: We deploy LaNE on tasks Franka Panda robot: reach fixed location, block, open a drawer, and insert a pen. The RL agent observes two RGB images, one in front of the robot and one on its wrist. Trained from scratch, the robot achieves a 10/10 evaluation success rate with only few demonstrations and less than one hour of learning.
5.3 Real Robot Experiments
As Figure 7 illustrates, LaNE enables efficient RL in four diverse tasks with the Franka Panda robot. From the easiest to the hardest, the tasks are Reaching Fixed Goal, Opening a Drawer, Lifting Block (random location), and Inserting Pen. Demonstrations are collected via teleoperation from a mobile app, which uses inside-out tracking to stream the device's pose.
Results show that our method extremely data-efficient, leading to task success with under hour of training. Table shows the number of demonstrations provided and training performance. For the most straightforward goal-reaching task, LaNE uses only one demo and completes the task for the first time in only 13 minutes. The agent finds consistently successful pol- icy quickly, only 8 minutes later. With five demonstrations, LaNE can learn more chal- lenging tasks with high stochasticity or re- quiring precision, including lifting block and inserting a pen. Overall, the results prove that our approach is outstanding in simulation and practical in the real world.
Task
Reach
Druwer
Lift
Inseri
Episodes Steps
Time
0:20
0:25
2:00
3:00
CoDER
Steps
N/A
N/A
N/A
N/A
First Success
Time
>> 1h
>> 1h
>> 1h
>>
LaNE (Ours)
Steps
476
1054
1109
1854
Time
LaNE (Ours) Convergenze
Steps
806
1929
1820
2927
Time
20:55
51:40
44:16
57:58
Table I: We deploy LaNE Franka Panda arm to manipulation tasks requiring up to 7 degrees freedom. Human demonstrations are provided via telcoperation, tak- ing only few minutes. LaNE trains every task 10/10 success rate from scratch under hour of wall clock time.
comparison, our strongest baseline CoDER [34] fails to succeed even once during the first hour of training.
Conclusion
This paper presents LaNE, data-efficient RL algorithm to learn sparse reward tasks from image observations by utilizing few demonstrations. LaNE outperforms state-of-ther thettr-rrtttttt benchmarks and enables under-an-hour RI. training with real robot. Our key innovation is to learn a latent dy- namics model, which provides a temporally consistent embedding space. When a transition lands sufficiently close to demonstration, we assign an extra task mask-prgrres-informee by the distance to the goal. As such, we convert sparse reward task to a task with dense proxy rewards, dramatically improving learning efficiency. Our work lends itself to exciting future direc- tions. For example, can leverage recent advancements in causal RL and counterfactual analysis [44, 45, 46] to determine the state in expert demonstration that directly caused task success. This might further improve our search for the nearest demonstration and overall learning.
reward modulated

---

Acknowledgements
This work was partly supported by ARL W911NF-21-1-0009, ONR 202 202102702AWD, DARPA ANSR: RTX CW2231110, and Lockheed Martin Corporation.
References
[1] T. Haarooja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, 2018.
[2] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski, et al. Human-level control through deep rein- forcement learning. nature, 518(7540):529-533 2015.
[3] Schulman, Wolski, Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 177.007777, 2017.
[4] J. Schulman. Trust region policy optimization. arXiv preprint arXiv:1502.05441 2015.
[5] O. M. Andrychowicz, Baker, M. Chociej, Jozefowicz. B. McGrew, J. Pachocki, A. Petron, M. Plappert, G. Powell, A. Ray, al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research, 39(1):3-20) 2020.
[6] D. Hafner, T. P. Lillicrap, M. Norouzi, and J. Ba. Mastering atari with discrete world models.
In International Conference on Learning Representations, 2021.
[7] M. Laskin, A. Srinivas, and P. Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, 2020.
[8] M. Laskin, K. Lee, A. Stooke, Pinto, P. Abbeel, and A. Srinivas. Reinforcement learning with augmented data, Advances in Neural Information Processing Systems, 2020.
[9] D. Yarats, I. Kostrikov, and R. Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations, 2021.
[10] G. Brockman, V. Cheung, L. Pettersson, J. Schneider, J. Schulman, J. Tang, and W. Zaremba. Openai gym, 2016.
[11] Y. Zhu, Wong, Mandlekar, and R. Martin-Martín. robosuite: modular simulation framework and benchmark for robot learning. In arXiv preprint arXiv:2009. 2020.
[12] B. Eysenbach, A. Gupta, J. Ibarz, and S. Levine. Diversity is all you need: Learning skills without a reward function. arXiv preprint arXiv:1802.0600 2018.
[13] K. Gregor, D. Rezende, and D. Wierstra. Variational intrinsic control. arXiv preprint arXiv.1611.07507. 2016.
[14] D. Kalashnikov, A. Irpan, P. Pastor, J. Ibarz, A. Herzog, Jang, D. Quillen, E. Holly, M. Kalakrishnan, V. Vanhoucke, al. Scalable deep reinforcement learning for vision-based robotic manipulation. Conference on Robot Learning, 2018.
[15) G. Cideron, B. Tabanpour, S. Curi, S. Girgin, L. Hussenot, G. Dulac-Amold, M. Geist, O. Pietquin, and R. Dadashi. Get back here: Robust imitation by return-to-distributiin plan- ning. arXiv preprint AXXi:2335.01100 2023.
[16] R. Dadashi, L. Hussenot, M. Geist, and O. Pietquin. Primal wasserstein imitation learning. arXiv preprint arXiv:2006. 04678, 2020.

---

[17] Y. Guo, J. Gao, Z. Wu, C. Shi, and J. Chen. Reinforcement learning with demonstrations from mismatched task under sparse reward. In Conference on Robot Learning, 2023.
[18] T. Hester, M. Vecerik, O. Pietquin, M. Lanctot, T. Schaul, B. Piot, D. Horgan, J. Quan, A. Sendonaris, I. Osband, et al. Deep q-learning from demonstrations. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
[19] M. Vecerik, T. Hester, J. Scholz, F. Wang. O. Pietquin, B. Piot, N. Heess, T. Rothörl, T. Lampe, and M. Riedmiller. Leveraging demonstrations for deep reinforcement learning on robotics problems with sparse rewards. arXiv preprint arXiv: 1707.08877 2017.
[20] Y. Wu, M. Mozifian, and F. Shkurti. Shaping rewards for reinforcement learning with imper- fect demonstrations using generative models. In International Conference Roboties and Automation, 2021.
[21] A. Nair, B. McGrew, M. Andrychowicz, W. Zaremba, and P. Abbeel. Overcoming exploration in reinforcement learning with demonstrations. In International Conference on Robotics and Automation, 2018.
[22] I.-C. A. Liu, S. Uppal, G. S. Sukhatme, J. J. Lim, Englert, and Y. Lee. Distilling motion planner augmented policies into visual control policies for robot manipulation. In Conference on Robot Learning. 2022.
[23] A. Rajeswaran, V. Kumar, A. Gupta, G. Vezzami, Schulman, E. Todorov, and S. Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstra- tions. In Robotics: Science and Systems, 2018.
[24] S. Haldar, V. Mathur, D. Yarats, and L. Pinto. Watch and match: Supercharging imitation with regularized optimal transport. In Conference on Robot Learning, 2023.
[25] S. Haldar, J. Pari, A. Rai, and L. Pinto. Teach a robot to fish: Versatile imitation from one minute of demonstrations. arXiv preprint arXiv:2303.01497, 2023.
[26] P. Ladosz, L. Weng, M. Kim, and H. Oh. Exploration in deep reinforcement learning: A survey. Information Fusion, 85:1 1-2, 2022.
[27] A. D. Laud. Theory and application of reward shaping in reinforcement learning. University of Illinois at Jrbana-Champaign, 2004.
[28] S. Cabi, S. Gómez, A. Novikov, K. Konyushova, S. Reed, R. Jeong, K. Zolna, Y. Aytar, D. Bud- den, M. Vecerik, O. Sushkov, D. Barker, J. Scholz, M. Denil, Freitas, and Z. Wang. Scaling data-driven roboties with reward sketching and batch reinforcement learning. In Robotics: Science and Systems, 2020.
[29] D. Pathak, P. Agrawal, A. A. Efros, and Darrell. Curiosity-driven exploration by self- supervised prediction. In International Conference on Machine Learning, 2017.
[30] D. Rengarajan, G. Vaidya, A. Sarvesh, D. Kalathil, and S. Shakkottai. Reinforcement learning with sparse rewards using guidance from offline demonstration. arXiv:2202.04622 04628, 2022.
arXiv preprint
[31] J. Ho and S. Ermon. Generative adversarial imitation learning. In Advances in Neural Infor- mation Processing Systems, 2016.
[32] J. Fu, K. Luo, and S. Levine. Learning robust rewards with adversarial inverse reinforcement learning. arXiv preprint 17Xiv:1710.1 1248, 2017.
[33] I. Kostrikov, K. K. Agrawal, D. Dwibedi, S. Levine, and J. Tompson. Discriminator-actor- critic: Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference Learning Representations, 2018.
10

---

[34] A. Zhan, R. Zhao, L. Pinto, P. Abbeel, and M. Laskin. Learning visual robotic control effi- ciently with contrastive pre-training and data augmentation. In International Conference on Robotics and Automation, 2022.
[35] N. Hansen, Y. Lin, H. Su, X. Wang, V. Kumar, and A. Rajeswaran. Modem: Accelerating visual model-based reinforcement learning with demonstrations, In International Conference on Learning Representations, 2023.
[36] D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Mastering visual continuous control: Improved auggeetuee reinforcement learning. arXiv preprint arXiv:2107.09645, 2021.
[37] D. Dwibedi, Y. Aytar, J. Tompson, P. Sermanet, and A. Zisserman. Temporal cycle-consistency learning. In Computer Vision and Pattern Recognition Conference, 2019.
[38] Sermanet, C. Lynch, Y. Chebotar, J. Hsu, E. Jang, S. Schaal, Levine, and G. Brain. Time- contrastive networks: Sele-super learning from video. In International Conference on Robotics and Automation, 2018.
[39] Z. Yuan, Z. Xue, B. Yuan, X. Wang, Y. Wu, Y. Gao, and H. Xu. Pre-trained image encoder for generalizable visual reinforcement learning. In Advances in Neural Information Processing Systems, 2022.
[40] M. Oquab, T. Darcet, T. Moutakanni, H. V. Vo, M. Szafraniec, V. Khalidov, P. Fermandez, D. HAZIZA, F. Massa, A. El-Nouby, M. Assran, N. Ballas, W. Galuba, R. Howes, P.-Y. Huang, S.-W. Li, I. Misra, M. Rabbat, V. Sharma, G. Synnaeve, H. Xu, H. Jegou, J. Mairal, P. Labatut, A. Joulin, and P. Bojanowski. DINOv2: Learning robust visual features without supervision. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
[41] M. Watter, Springenberg, J. Boedecker, and M. Riedmiller. Embed control: A locally linear latent dynamics model for control from raw images. In Advances in Neural Information Processing Systems, 2015.
[42] D. P. Kingma and M. Welling. Auto-Encoding Variational Bayes. In International Conference on Learning Representations, 2014.
[43] H. Van Hasselt, A. Guez, and D. Silver. Deep reinforcement learning with double q-learning. In Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
[44] E. Bareinboim and Pearl. Causal inference and the data-fusion problem. Proceedings of the National Academy of Sciences, 113(27):7345-7332 2016.
[45] T. Mesnard, T. Weber, F. Viola, S. Thakoor, A. Saade, A. Harutyunyan, W. Dabney, T. Steple- ton, N. Heess, A. Guez, et al. Counterfactual credit assignment in model-free reinforcement learning. arXiv preprint arXiv:2011.09464, 2020.
[46] S. Zhu, I. Ng, and Z. Chen. Causal discovery with reinforcement learning. arXiv preprint 190.:1906.04477.. 2019.
[47] J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
[48] X. Chen, S. Xie, and K. He. An empirical study of training self-supervised vision transformers. In International Conference Computer Vision, 2021.
[49] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
[50] A. Albu-Schaffer, C. Ott, U. Frese, and G. Hirzinger. Cartesian impedance control of redundant robots: Recent results with the niight-weight-anne In International Conference Roboties and Automation, 2003.
11

---

[51] S. R. Buss. Introduction to inverse kinematics with jacobian transpose, pseudoinverse and damped least squares methods. IEEE Journal of Robotics and Automation, 17(1-19):16, 2004.
[52] J. Linowes and K. Babilinski. Augmented reality for developers: Build practical augmented reality applications with unity, ARCore, ARKit, and Vuforia. Packt Publishing Ltd, 2017.
12

---

7 Appendix
7.1 Motivating Example
We motivate the need to learn a dynamics-aware embedding space: finding a good distance measure between states from their respective high-dimensional image observations non-trivial. We use the PointMaze environment from the D4RL benchmark [47] to provide a clear illustration. In this environment, the controllable point is marked in green. As shown in Figure we place the point at three positions in the maze such that state (a) is much closer state (b) than state (c).
E
= (b) (-0.5, = (1,1)
x
Figure 8: Observations from the PointMaze environment. The point mass in green the controllable agent, whose location indicated by s. indicated by the agent's location, state b is closer state than
Distance Measure
d(a, b)
d(a, c)
d(a, c)/d(a,
Pixel L
2.835
3.303
1.165
DINOv2 T.2
0.004
0.003
0.75
Ground Truth L
0.5
1.5
ADM (Ours)
0.046
0.135
2.935
Table 2: Comparison across different distance measures PointMaze. Distance pixel space NOv2 emhedding space cannot capture the true relations of the underlying states. Whereas our proposed ADM can estimate the relative distance the most accurately.
le
1e-5
1.5
3
3
1.0
2
2
1
-1
0.0
Figure 9: Left: ADM (Ours). Middle: DINOv2. Right: ADM without DINO. We plot the L2 distance ADM and DINOv2 embedding space from each location the maze the pink location. In the ADM embedding space, computed distances match the environment dynamics, showing the lower left corner as the farthest due to the turn. In contrast, DINOv2 cannot identify the differences between states. Finally, ADM is trained without DINOv2, it suffers from poor convergence due to difficulty in feature learning.
Table compares two baseline distance measures with our proposed method ADM, trained from 10000 random interactions in the environment. For Pixel 1,2 directly compute the Euclidean distance. As expected, d(a, b) and d(a,c) not distinguishable from the pixel-wise distance. Similarly, the pre-trained DINOv2 embeddings do not capture the transition information and fail to estimate the distance between states. ADM is the only method to capture the relative proportions of the ground-truth distances.
We visualize the quality of the learned embeddings in Figure where we plot the estimated distance from all locations within the maze fixed location. We show that the embeddings learned by LaNE respect the dynamics of the maze, understanding that the point must go around the corner instead of
13

---

3
1.0
2
0.5
1 0
0.0
0
-1
Figure 10: Left: ADM (Ours). Middle: DINOv2. Right: ADM without DINO, Instead of L distance, we plot d(v1, 02) := cos(U), 02), the negative cosine distance between the embedding vectors. The scale 10-7 indicates that these vectors all point in similar directions. Hence, cosine similarity is not as informative as L,?.
through the wall. This simple experiment illustrates the benefits of using our latent space distance measure to quantify the task-relevant similarity between image observations.
Additionally, we are curious if cosine similarity, commonly used for vector retrieval, could be used in place of 2 distance. As illustrated in Figure 10, all embedding vectors point in similar directions, making the cosine similarity less informative.
7.2 Visualization of Embedded Robot Trajectories
This section aims to provide better intuitive understanding of the learned ADM embedding space. We visualize the learned embeddings in the real robot reach task by plotting their 2-D projections in Figure 11. Specifically, we plot demonstration, random successful, and random unsuccess- ful episode. We observe a clear separation between successful and failed episodes: the successful episode closely follows the demonstration. In addition, we observe that the latent states move lin- early as the agent progresses in the task.
Demo
10
Success
-10
-20
Fail
-30
-20
Figure 11: 2-D t-SNE of a demonstration, successful and an unsuccessful trajectory on the Reach task. Colors change from light dark as the episodes progress. The successful episode is mapped closer to the demonstration than the failed episode.
7.3 Comparison with Imitation Learning Methods
In this section, compare LaNE with powerful imitation learning algorithms, namely FISH [25] and PWIL [16]. FISH builds non-parametric base policy and trains an RL residual policy to minimize the earth mover's distance between the policy rollout and demonstration distribution. Similarly, PWIL uses RL to minimize an upper bound of the Wasserstein distance between rollouts and demonstrations. FISH uses an encoder pre-trained with behavior cloning as a feature extractor, whereas PWIL uses TCC [37] model pretrained on demonstration trajectories.
As shown in Figure 12, both IL baselines fail to solve the Robosuite tasks given the same number of demonstrations LaNE. One potential cause is that their feature extractors learned solely from
14

---

demonstrations and hence struggle to generalize to OOD observations, which often occurs when the 7-DoF end-effector is allowed to rotate. On the other hand, these II. methods don't utilize any environment signals during RL training. Future works can explore the combination of sparse task rewards, online representation learning, and IL methods.
PNL
LaKC
L
Doar
1.
1.0
1.0
0.8
0.8
0.5
0.6
0.8
0.4
02
0.
0.9
022
044
0.5
08
0.0
0.5
3.0
05
1.5
2.0
25
renvrnmment uteps
steps
staps
#entrrnmeett stepa
Figure 12: We compare LaNE with two imitation learning methods. Since the IL methods don't utilize the sparse environment rewards, they fail learn the tasks from the limited number of demonstrations.
7.4 Other Ablation Studies
To justify the design choice of LaNE's exploration reward shown in Equation 13, we conduct ab- lation studies on variants of the exploration reward. 1. We set necessary. Alternatively, we use simpler reward based the L2 distance between the ob- servation and its nearest-neighhor demonstration: re max(1, o(o This simplified reward is the inverse 1.2, scaled to the average distance between demonstration states, and clipped to make the g-value bound hold. Results on the left of Figure 13 show that our method outperforms both alternatives and holds a clear advantage the long-horizon can-moving task.
verify the discounting
Aside from our Augmentation-IIvarian Distance Measure, we experiment with other ways to learn image embeddings. 1. We verify the necessity of the local linearity constraint by swapping the locally linear dynamics model with fully connected neural network. 2. We perform contrastive learning with MoCo-v3 [48] on individual images without learning the dynamics. The right two plots in Figure 13 demonstrate that our method is advantageous over the alternative representation learning approaches.
1.0
1.0
0.8
0.8
550000
0.4
0.0
0.0
0.0
10
2.5
0.5 envermment steps 145
0.5
1.5
emu memee steps
PaTARrEmeT daen 145
Rensitmnn steps
We compare two other choices of exploration reward Te- bonus is not discounted based on the estimated distance success. L2 inverse reward is simple reward based on a state's 's estimated distance to its closest demo. Both alternatives are inferior to LaNE. Right: We compare two other representation learning methods to construct distance measure. Non-Linear: We don't enforce local linearity in the forward model and use fully connected network. MoCo-v3: We don't train forward model use contrastive learning on individual images. LaNE holds clear advantage in the block stacking task. In move-a-can, LaNE the earliest to achieve success.
Figure 13: Left:
indicates the exploration
7.5 Comparison with Dreamer
We compare LaNE with DreamerV2, a state-of-the-ar model-based RL method. Because Dream- e2V2 does not explicitly take demonstrations, we pre-fill its replay buffer with demonstration trajec- tories. In addition, we allow DreamerV2 to run for longer to get a better sense of its performance. As shown in Figure 14 below, it takes over five times as many environment steps for DreamerV2 to learn the Robosuite block-lifting task compared to LaNE. We note that DreamerV2 does not dis- tinguish the pre-filled demonstrations from the regular episodes and thus fails to sample the sparse
15

---

reward often enough. Nevertheless, LaNE can also integrate with model mosel- RL methods to boost their sample efficiency in sparse-reward settings.
1.0
DreamerV2 LaNE
Success Rate
0.8
0.4
0.2
0.0
1.0
1.5
2.0
2.5
environment steps
Figure 14:
Dreamer takes over 5 times samples to learn the Robosuite Lift a Block task compared to LaNE because it docs not distinguish between demonstration and regular episodes.
Robustness LaNE Policies
We evaluate the out-of-distribution robustness of trained LaNE policies. Specifically, push the block and the pen with stick during policy execution, illustrated below in Figure 15. We find that the policies are robust to small perturbations. For example, the agent can lift the yellow block when pushed, even when the stick and the operator's 's arm enter the camera view. Additionally, when the pen is moved to locations unseen during training, the robot can still pick it up.
Figure 15:
The learned LaNE policies are robust to small perturbations during evaluation, including moving objects, unseen initialization locations, and foreign objects.
7.7 Environment Details
Maximum Episode Length: For our simulated experiments Robosuite, set maximum episode lengths based on the difficulty of each task, shown in Table 3. These numbers are slightly over the average steps the demonstrator takes to complete each task, leaving the RL agent plenty of time to finish. For our real robot experiments, all four tasks share the same maximum episode length of 30 steps. The episodes are terminated if the maximum length is reached or when the task is completed.
Task
Lift
Door
Stack
Move Can
Max Episode Length
40
80
80
120
Table We choose the maximum episode length of each task slightly over the average steps the demonstrator took, giving the RL agent ample time to finish.
7.8
Real Robot Setup
7.8.1
Highly Compliant Real-time Controller
Because reinforcement learning requires a trial-and-error process [49], we expect the robot make frequent physical contact with the environment. To eliminate safety hazards, the controller must be compliant with external forces. On the other hand, want the robot to move swiftly that each
16

---

RL step takes less time to execute. We develop highly compliant real-time controller extended from the Cartesian Impedance Controller [50, 51].
At a high level, the end-effector tracks an equilibrium pose following a ceeee-sprin-daaaaap model. As the current end-effector pose deviates further from equilibrium, the robot asserts higher torque in the opposite direction. In the Cartesian space, we limit the maximum force exerted on the end- effector by the robot, preventing it from causing damage. In the joint space, we apply a counter torque when joint gets close to its hardware limit. Combining these control rules builds a safety net around the robot for smoother RL training.
7.8.2 System Architecture
The robot uses two Intel Realsense cameras, one mounted on the end-effector and another in front of the robot. Our method only utilizes color images. Both cameras are connected via USB to an Nvidia GPU desktop, which runs inference and training for the RL agent. Specifically, the GPU desktop runs OpenAl gym interface [10], with which the RL agent interacts. At each timestep, the agent chooses its action based on its policy: effector position, change in roll/pitch/yaw, and open/close of the gripper. Next, the action selected by the RI. agent is sent via ethernet to an Intel NUC, which directly interfaces with the Panda Robot. Specifically, the Intel NUC runs the Robot Operating System (ROS), where our real-time controller communicates with the hardware.
((). The action consists of a displacement of the end
7.8.3 Demonstration Collection via Tele-Operation
A key aspect of our approach is learn from small set of hu- man demonstrations efficiently. We build an application where user can tele-operate the robot by moving and rotating smartphone. Our iPhone application utilizes primitives from Apple's ARKit [52] to stream the position and orientation of the device to the PC controlling the robot.
XR Controller
During demonstration collection, Python script translates the tracking data into gym actions and executes them on the real robot. The gym environment updates the robot's equilibrium pose to follow the phone's 's movement. These demonstration trajectories are stored on the Nvidia GPU desktop and used during training.
Settins
Recovry
II
7.9 Hyper-parameterr and Training Scheduling
We use discounting factor 7 0.99 and exploration ward discount a 0.98 for all our experiments. The actor and critic learning rates are 10-3, The latent dynamics model learning rate is 4- 10-3 We use a training batch size of 128. We choose different fractions of demonstrations Pa when we perform prioritized replay, as shown in the table below.
We develop a mobile app to stream the device pose to the PC. This allows us to collect demonstra- tions the real robot arm quickly.
Figure 16:
Task
Lift
Door
Stack
Move a Can
Real Rohot
Pd
0.15
0.15
0.15
0.2
0.2
Table 4: We vary the fraction pa of demonstrations within each batch depending on the task we are training.
For simulated environments, we perform one Actor-Critic update for each environment step. After loading the demonstrations, we update the latent dynamics model every 5000 environment steps until convergence. For real robot experiments, we update the latent dynamics model every 30 steps the robot takes. Since each step takes about 0.5 seconds to execute the real robot, we keep performing SAC updates in the background while the robot is in motion.
17

---

7.10 Neural Network Architectures
Our method consists of 6 components parameterized by neural networks, namely: model encoder Eg. model decoder Dg. locally-linear dynamics model Mø. RL encoder ËRL, Actor TRL and Critic QRL- The numbers below correspond to our specific setting where the latent space has 16 dimen- sions, and the action space has 7 dimensions. Our locally linear dynamics model predicts low-rank approximation of 16 16 matrix A using two 16-dimensional vectors and where A I + uv".
Model Eacoder (vith DINDv2):
Input: 3x112x112 randonly crapped inages
Pretrained DINOV2 varient: dinov2.Pita14_
Concatenate of DIB0v2 embeddinge
RaLU (L.nnarr (ont_featurea==513))
(Linear (out, (out_featurea-s12)
Linear (out_ saut__ffature=322
Output:
nean
log-std
Model Decoder (vith DINDv2):
Input: 16-din
vectors
RaLu (ot__feturure==511222 (Linear RaLU (Linear (out_teaturea==813) Linear fout__fetuueeeee 768)
Output: predicted DINDv2 embeddinge
Modal Encadar (Ne DINOv2):
randonly cropped inages Input: (CoavD (6, 32, kernel-3, stride-2))) (Lyyeelayrl
RoLU (Layertorn (Conv2D (32, 32, kernel-3))
ReLU (LayerWorn Platten
32, kernel-3))
RaLll (Linear (out_ features==18))
RetU (ottterrrtttttttureeeete
(out_feee featuree-32)
Output:
16-din mean 16-din log-std
Model Decoder (No DINOV2): Input: 16-din latent vectors
ReLU (Linear (ottaaar ReLU (outteerroe___etueeeeee
ReLU (Lutneirrruut_feeeetee Reshape (out_feature==322785
128x16x16
Upsample into 128x32x32
(eLU(Co
kernel-3, stride=1, pad-1)))
Upsample into 128x64x64
(eLU(Co (128, 128, xernel-3, stride-1, pad-1)))
Upsemple into 128x126x128 Coav2D 6, kernel=3, stride=1, pad-1))
Output: 6x128x128 inages
Loieery-Linear Dynamics model: Input: 16-din latent Raill (L.tnear(ot__featuree
[Linearrout_ffaaureeeetttt
RaLl RaLU (oatarootttttteeaatttaee
Dutput: 16-din vector 16-din vector
16x7 natrix 16-din offset
Encoder:
Input: 6x112x112 randonly cropped inagea ReLU (Layerror (Coav2D 32, kernel=3, stride-2)))
(LayerNorn (Conv2D (32, kernel=3, stride=2)))
RetV (LayerLee (Coav2D (32, 32, kermel-3, stride-2)))
32, kornel=3, stride=2)))
Flatten
LayerNorm (Linear (attureaat))))))))
Output: 32-din feature
Input: 32-din feature
RaLU (Linear (ottrrootttttuee-11000. RaLU (Linear f
Linear (out_festtt Output: festuree-14)
7-din log-atd; Critic: Q-value
18


Reviews:
[{"review_id": "2VFj0k9stu", "reviewer": "Reviewer_ic1h", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "In this paper, the authors present a visual RL method called LaNE, designed to efficiently learn a policy for robotic manipulation tasks with sparse rewards from image observations and a few demonstrations. LaNE achieves superior sample efficiency by learning an embedding space to quantify state proximity and reward explorations near the demonstrations. Experiments demonstrate that LaNE can train a real Franka Panda robot from scratch within an hour using only a few demonstrations.", "summary_of_recommendation": "The paper is well-written and presents a novel visual RL method, LaNE, that efficiently learns policies for robotic manipulation tasks with sparse rewards using image observations and limited demonstrations. However, several issues need to be addressed, including providing additional ablation studies comparing linear and non-linear dynamics models, experimenting with task-progress-informed exploration rewards without sparse rewards, and addressing discrepancies in the performance of MoDem. Additionally, open-sourcing the code and providing visualizations of the representation for $z$ with and without DINO would greatly enhance the reproducibility and robustness of the results. Addressing these issues would significantly strengthen the contribution and impact of the work.", "technical_quality": 3, "clarity_of_presentation": 3, "potential_impact": 2, "robotics_focus": 4, "quality_of_limitations": 2, "questions_for_rebuttal": "Issue 1: The authors train a latent dynamics model in a linear system manner. Why not train a neural network for the dynamics model, which can handle non-linearity? Please provide an additional ablation study comparing the linear dynamics model to a non-linear dynamics model.\n\nIssue 2: Using sparse rewards remains an issue in real robot environments. Please provide additional experiments using only task-progress-informed exploration rewards without relying on sparse rewards.\n\nIssue 3: In Fig. 5, the performance of MoDem presented by the authors is significantly worse compared to the performance reported in the MoDem paper. To objectively evaluate the performance comparison, we request the authors to open source their code. The authors can also perform a comparative evaluation with MoDem-v2 in a real robot environment.\n\nIssue 4: In Fig. 6, the authors present ablation studies on adding PR (prioritized replay), VC (value clipping), and exploration reward $r_e$, based on RAD. However, some combinations are missing, such as RAD+VC and RAD+$r_e$. Since the effect of $r_e$ is expected to be the most pronounced, we request additional ablation studies for each combination.\n\nIssue 5: In Fig. 6, there is little performance difference between w/o DINO and w/ DINO, which may suggest that a CNN-based encoder and decoder are sufficient without using DINO. To confirm this, a visualization of the representation for $z$ is needed for both w/o DINO and w/ DINO conditions.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NTPqhV9EBP", "reviewer": "Reviewer_9FLB", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "Train a dynamic aware autoencoder, and using nearest neighbor in its latent space, give extra reward for exploration", "summary_of_recommendation": "Given the rebuttal, I updated my assessment to a weak accept", "technical_quality": 3, "clarity_of_presentation": 3, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 2, "questions_for_rebuttal": "1. Why compare with non-demo guided RL methods if there are so many good demo-guided RL methods out there? This seems like an obvious oversight, what am I missing?\n2. What happens if we use simple image-space encoder fine-tuning like MoCo-v3?\n3. Can this pretraining improve, say, imitation learning methods?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nlDWOQbArU", "reviewer": "Reviewer_qY4v", "summary": "", "questions": "", "limitations": "", "rating": 2, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This work introduces a novel approach to designing dense rewards for Learning from Demonstration (LfD) tasks. The method, named LaNE, employs a locally linear dynamic VAE to encode features and generate a feature measurement metric, leveraging the vision foundation model DINOv2. During skill learning with Reinforcement Learning (RL), LaNE provides an auxiliary reward when the current image observation is deemed similar to the demonstrations. The authors assert that this augmented reward design significantly accelerates LfD tasks, offering a new perspective on bridging the gap between sparse demonstration data and efficient skill acquisition in robotics.", "summary_of_recommendation": "This work is clearly written, despite it shows the nearest retrieval augmentation can accelerate skill learning, it still lacks some related work comparisons and makes the conclusion less convincing..", "technical_quality": 2, "clarity_of_presentation": 2, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "Q1: Comparison with Existing Dense Reward Methods\n\nThis work aims to accelerate Learning from Demonstration (LfD) with Reinforcement Learning (RL) by creating an augmented dense reward. However, designing dense rewards for LfD isn't always problematic. Previous works such as Temporal Cycle Consistency (TCC) and Time Contrastive Networks (TCN) have demonstrated the ability to learn stage representations of tasks and provide dense rewards, accelerating sparse reward LfD tasks. How does this work compare to or potentially complement these existing approaches? It would be valuable to see a comparative analysis or exploration of combining these methods to conclusively determine if the nearest retrieval augmented reward can significantly boost LfD in sparse reward settings.\n\nQ2: Comparison with Existing Imitation Learning Reward Functions\n\nThe concept of designing a demonstration-guided reward is intriguing. However, there are existing imitation learning reward functions designed to accelerate LfD tasks. A notable example is Primal Wasserstein Imitation Learning (PWIL), which designs a dense reward by calculating greedy Wasserstein distance and has proven effective even when learning from a single demonstration. PWIL also evaluated using L2 distance as a reward to guide imitation. How does the proposed method compare to or differ from these existing approaches, particularly PWIL?\n\nQ3: Role of Representation Learning Module\n\nIf the representation learning module (local linear dynamic VAE) is key to LaNE's effectiveness, it would be valuable to combine it with existing methods like PWIL. This combination could verify whether the local linear dynamic VAE plays a vital role in accelerating LfD across different reward design approaches. Have the authors considered or conducted such experiments?\n\nQ4: Necessity of DINOv2\n\nIn Figure 6 (right), the ablation study seems to suggest that DINOv2 might not be necessary for LaNE. Could the authors elaborate on this finding? Additionally, for finding similar images, using cosine similarity instead of L2 distance might be more appropriate, as it's often used in deep image retrieval and may perform better for high-dimensional features. Have the authors considered this alternative metric?\n\nQ5: Impact of Demonstration Quality and Multimodality\n\nThe augmented reward appears to constrain the policy to explore more around the demonstration region. This raises an interesting question: how do the quality and multimodality of demonstrations affect the policy learning performance in LaNE? Have the authors explored the robustness of their method to variations in demonstration quality or the presence of multiple valid ways to complete a task?\n\nQ6: Discounting the exploration reward based on estimated distance to the goal\n\nIt is better to have ablation study to check how such design can accelerate/harm the learning.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Strengths
- Interesting method
- Nice evaluation in simulation and on a real robot

Weaknesses
- Doubts about limited setting (sparse rewards)
- Missing details and justifications
- Doubts about results of some of the comparisons
- Baselines don't seem to be the most relevant ones

# After Discussion Phase
Both reviewers ic1h and 9FLB were very happy with the replies and additional material/experiments and accordingly updated their scores. Reviewer ic1h still believes that adding a feature visualization to confirm the effectiveness of using DINOv2 would make the paper stronger. Reviewer qY4v unfortunately did not react, but also those comments were well addressed and the requested experiments and ablations added (with results further supporting the method).
Decision: Accept