Student-Informed Teacher Training
Authors: Nico Messikommer
Jiaxu Xing
Elie Aljalbout
Davide Scaramuzza
Venue: ICLR 2025 Spotlight (2025)
Keywords: Reinforcement Learning
Imitation Learning
Robotics
Primary area: reinforcement learning
TL;DR: To address the teacher-student asymmetry in imitation learning, we propose a joint learning framework for both teacher and student, adapting the teacher to the capabilities of the student during training to enhance alignment.

Abstract: Imitation learning with a privileged teacher has proven effective for learning complex control behaviors from high-dimensional inputs, such as images. In this framework, a teacher is trained with privileged task information, while a student tries to predict the actions of the teacher with more limited observations, e.g., in a robot navigation task, the teacher might have access to distances to nearby obstacles, while the student only receives visual observations of the scene. However, privileged imitation learning faces a key challenge: the student might be unable to imitate the teacher's behavior due to partial observability. This problem arises because the teacher is trained without considering if the student is capable of imitating the learned behavior. To address this teacher-student asymmetry, we propose a framework for joint training of the teacher and student policies, encouraging the teacher to learn behaviors that can be imitated by the student despite the latters' limited access to information and its partial observability. Based on the performance bound in imitation learning, we add (i) the approximated action difference between teacher and student as a penalty term to the reward function of the teacher, and (ii) a supervised teacher-student alignment step. We motivate our method with a maze navigation task and demonstrate its effectiveness on complex vision-based quadrotor flight and manipulation tasks.

Published as a conference paper at ICLR 2025
STUDENT-INFORMED TEACHER TRAINING
Nico Messikommer's Jiixx Xing", Elie Aljalbout, Davide Scaramuzza Robotics and Perception Group, University of Zurich, Switzerland 'Equal contribution
Overview
Roll Out Phase
Policy Update Phase
Aligrment Plase
Teacher Obs
Student
Teacher Olte
Teacher
Teacher Ott
Teacher
Teacher Obs
Student Obs
Fr
f
R
Founy
R
F
E3-Dic Lak:
LILLON
111ss
KI- Penaliyy
KI- Loss
Paind L1-1.oss
(a)
(b)
(c)
(d)
Figure I: Method Overview. (a) We train three networks by freezing weights (grey box) and chang- ing gradient flows (dashed arrow) in alternating phases.
(b) In the roll-out phase, the KL-Divergence between the proxy student Fs and teacher Fr is used as penalty term. (c) Additionally to the policy gradient, the teacher encoder updated by backpropagating through the KL.-Divergence between the action distribution of the teacher and the proxy student. (d) Using student observations, the proxy student is aligned to the student E's and the student to the teacher network.
ABSTRACT
Imitation learning with a privileged teacher has proven effective for learning com- plex control behaviors from high-dimensiooal inputs, such as images. In this framework, a teacher trained with privileged task information, while a student tries to predict the actions of the teacher with more limited observations, e.g.. in a robot navigation task, the teacher might have access to distances to nearby obstacles, while the student only receives visual observations of the scene. How- ever, privileged imitation learning faces key challenge: the student might be unable imitate the teacher's behavior due to partial observability. This problem arises because the teacher is trained without considering if the student is capable of imitating the learned behavior. To address this teacher-student asymmetry, we propose framework for joint training of the teacher and student policies, encour- aging the teacher to learn behaviors that can be imitated by the student despite the latters' limited access to information and its partial observability. Based on the performance bound in imitation learning, we add (i) the approximated action difference between teacher and student as a penalty term to the reward function of the teacher, and (ii) supervised teacher-student alignment step. We motivate our navigation task and demonstrate its effectiveness complex
method with vision-based quadrotor flight and manipulation tasks.
Multimedia Material The project website is httph::t////rt ifi uzh. ch/sitt/
INTRODUCTION
In reinforcement learning (RL), agent learns to perform task by interacting with its environment and maximizing the cumulative rewards gained through these interactions. RL has been shown outperform human abilities in several domains (Vinyals et al., 2019; Mnih et al., 2015; Silver
This work was supported by the European Research Council (ERC) under grant agreement No. 864042 (AGILEFLIIGHT)

---

Published as a conference paper at ICLR 2025
et al., 2016; Kaufmann et al., 2023). However, this process requires extensive exploration, as the agent must avoid getting trapped in local minima, often resulting in a large number of environment interactions (Pathak et al., 2017). The number of interactions is even further increased when the agent processes high-dimensional data as input (Ota et al., 2020). Using such observations, the policy must learn to extract a notion of the agent's 's state, process that is computationally expensive when optimized soiely through RL. Improving the efficiency of RL training from high-dimensional observations can unlock significant advancements in various applications, particularly for robots interacting with the real world.
State-of-thh onf-the-art art (SotA) approaches often rely on imitation learning to accelerate training (Chi al., 2023; Luo et al., 2024; Team et al., 2024; Doshi et al., 2024). However, collecting expert demon- strations for imitation learning can prohibitively expensive, which has led to the development of the teacher-student framework. In this framework, expert data is generated automatically by training a teacher policy using RL on privileged task information, benefiting from efficient simulation and a faster learning process. For instance, if we take the example of autonomous driving as considered in (Chen et al., 2019), the teacher can have access to the layout of the environment, and the posi- tions of all traffic participants as part of its observation space. In the second phase, student policy, processing digh-dimensional observations as input, is trained by using the actions of the teacher as direct supervision signal. In the autonomous driving scenario, the student would attempt to infer the teacher's 's actions from visual observations. This approach eliminates the need for the student to extensively explore the environment, which can be a very challenging process when dealing with high-dimensions observations, such as images. By leveraging the direct supervision of the teacher, imitation learning significantly speeds up the training process for the student.
However, privileged imitation learning can be hindered by information asymmetry between the teacher and student, where the student receives less informative observations and struggles to imitate the behavior of the teacher (Nguyen et al., 2023). As consequence of the information asymme- try, the teacher tends to over-rely on its full observability of the environment without considering the more limited observation space of the student. This causes the teacher to provide target actions that the student cannot infer from its observations, since the student lacks access to the same level of environmental information. For example, consider a mobile robot navigating obstacle-fili fflllee environment. In this example, an information asymmetry in the observation space could easily arise if the teacher policy receives the relative distances to all surrounding obstacles, while the student is limited to a forward-face facing camera, requiring that obstacles be within the view of the camera.
To tackle this challenge, we propose teacher-student knowledge distillation framework that en- courages the teacher to learn behaviors that account for the capabilities of the student. Specifically, the objective function of the teacher is extended by adding the upper bound of the student perfor- mance within the imitation learning setting. This results in a reward term that penalizes the teacher for visiting states where there is a significant action mismatch between the student and teacher. Ad- ditionally, minimizing this upper bound leads to second optimization term that directly supervises the weights of the teacher network. A key feature of our approach is the dynamic prediction of the student capabilities during the teacher-environment interactions, which eliminates the need to directly render high- high-dimensional student observations. By embedding the upper bound of student performance for imitation learning into the learning process of the teacher, we effectively account for the teacher-student information asymmetry.
To validate our approach, we first test it in maze setting where the teacher can choose shortcut that is invisible to the student. Our method successfully adjusts the behavior of the teacher to take sub-optimal route, accounting for the limited observability of the student. Additionally, we apply our framework for training a robot manipulator to open a drawer while minimizing self-occlusion in front of camera. In this scenario, the teacher has access to the relative pose of the cabinet in addition to the internal state of the manipulator, while the student needs to compute the position of the drawer based on an image. Our method enables the teacher to modify its behavior reduce self-occlusion, leading to camera-aware behavior. Finally, we demonstrate the effectiveness of our method in the complex task of quadrotor flight through obstacles. The teacher, using state- based information, learns to orient the camera forward to help the student detect obstacles. Overall, our method leads to substantially higher student returns, reducing the gap between the teacher and student across tasks. As a result, our approach leads to very notable improvements in the student success rate in all considered environments.
2

---

Published as a conference paper at ICLR 2025
2 BACKGROUND
The goal of Reinforcement Learning (RL) is to find a policy that optimizes the expected return, con- sisting of the accumulated and weighted reward terms obtained by trajectories governed by the un- derlying Markov Decision Process (MDP). The MDP is defined as a tuple M = (S,A, P. R, y, Ho) comprising of the state space S and the action space A. At the start of episode, the agent is initialized based on the initial state distribution dn and outputs for each timestep an action or E A sampled from the agent policy <--ss). The transition to the state $1+1 follows the transition probabil- ity P(s++1|s;, az). The performance of an agent solving the MDP can be quantified by the expected
sum of rewards discounted by y: ((T) = Re IEi= 7rr(st, ax) 80 ~ μoo 0t ~ =T-1ss), 8++1 ~ P(-|st, 0t)]. By iteratively interacting with the environment, RL. algorithms try to find the optimal policies that maximize the expected return x arg max; J(x).
In contrast, imitation learning starts already with expert policy тT. which is imitated by the student policy Ts. This expert policy either directly accessible to the agent or is assumed to be previously used to collect an offline dataset of expert interactions with the environment. This setting eliminates the need for extensive exploration through environment interactions. Furthermore, a stricter supervision signal can be employed by directly providing ground truth actions from the expert to the student. However, if the data distribution for training the student policies does not represent the final application distribution, the student policy learned with imitation learning will suffer since it did not learn how to behave outside the training data distribution. In general, Xu et al. (2020); Syed & Schapire (2010); Ross al. (2011) established an upper bound on the difference between teacher performance Jur and student performance Jxs assuming the reward function is bounded by Fmax, i.e. [r(s, a) S Tmaxs
((TT) (TT)) S 2//2 max Ve,
(1 7)-
where e represents the upper bound of the expected action difference between student and teacher under the discounted stationary state distribution of the expert dxy (s) (1 - 7) Ei PrPr(se 8; TE)
Bxndez
(s))))ssss(s))())) as(( S
Imitation learning focused on minimizing this KL-Divergence by updating the student policy Iт to predict the same action distribution as the teacher TT-
3 STUDENT-INFORMED TEACHER TRAINING
As shown in Eq. 1 and Eq. 2, the performance gap between student and teacher is upper-bounded by the action difference between both policies. Thus, minimizing the action difference under the state distribution of the expert also minimizes the performance gap J(Tr) JTr). Instead of trying to minimize the action difference by adjusting the student policy πy (Nguyen et al., 2023; Shenfeld et al., 2023; Walsman et al., 2022), we propose to change the perspective and find teacher policy лт optimizing for the task reward while considering the alignment between teacher and student. Consequently, our framework requires that the teacher and student can be jointly trained. We want to find a teacher policy गT that maximizes the following objective
J(xT) Exnd, p ay (-19) [T(ss a)] Barde,,
(3)
Egdde a aw (t)) (s) DKI (=1s),
(4)
x ETrpp (R(T) Do(T)) (00==/oo00))))) po(T) (o(T))dr.
(5)
In the last step, the discounted state distribution is changed to the expectation over trajectories T ~ pa. which induced by the expert policy गт and represent the state and corresponding actions
T = {So, au, $1, a.....) Additionally, we define the return R(T) = Empaeer 7'r(s,, az) and the sum of discounted KL-Divergences Do(T) E.,er r¹ DD ( (TT(()))))
(=s() We use subscript 8, to emphasize that the probability distribution over the trajectories pe and Do is dependent on the parameter of the teacher network 6. Following the classical policy gradient to obtain the optimal policy, we take the gradient of Eq. 5 with respect the teacher parameters 0
3

---

Published as a conference paper at ICLR 2025
VOJ(TT) Ve (o(+)(((()) Do(T))dT
(6)
Vopo(T) (opo(t)Do(t)dr -ooo-pooooooaaooo po(o))ooo(tt)d
(7)
Vopo(T)(R(()) -Do(()) po(r))oDo(t)dr
(8)
Policy Gradient
KL-Div Gradient
As can be observed, we end up with the standard policy gradient optimizing the task reward while also considering the teacher-student misalignment for each trajectory. This first KL-Divergence D can be interpreted as a reward encouraging the teacher policy to visit states where the student and teacher are aligned and avoid states with a large misalignment. The additional "KL-Div Gradient" term contains the second KL-Divergence and represents the expectation of the gradient with respect to the teacher network over the expert states, which represents a direct supervision on the teacher weights by enforcing the prediction of the same action distribution as the student.
4 JOINT LEARNING FRAMEWORK
Building on the formulation in Sec. 3, we propose a practical framework to tackle the teacher-student asymmetry. Following Eq. 8, we adapt the widely-used PPO algorithm (Schulman et al., 2017) to train the teacher to learn behaviors that can be imitated by the student. At the same time, we train the student network to imitate the teacher based on subset of the collected environment interactions containing student and teacher observations. By using a subset of paired teacher and student data, avoid the (usually expensive) simulation of student observation for each time step the teacher interacts with the environment. An overview of the proposed method is shown in Figure a) and in the Appendix A.2. In the following sections, we introduce the policy networks employed in framework (Sec. 4. and detail the different training phases (Sec. 4.2).
4.1 ARCHITECTURE
To implement the objective in Eq. inside the teacher training, we implement two key components: (i proxy student network taking as input teacher observations and (ii) a shared action decoder network. Excluding the critic, our method consists of three different networks: the teacher network Fr, the student network Fs. and the proxy student network Fs. which all share the same action decoder network A.
Proxy Student Network To compute the action difference used in the penalty term and to obtain the KL.-Div Gradient in Eq. 8, a forward pass through both the teacher and student networks is required for each collected sample. However, simulating high-dimensional student observations, such as images, is often computationally expensive, contradicting the initial goal of accelerating training. To avoid this simulation overhead, we introduce a separate neural network Fs. Specifically, given the teacher observations, the proxy student Fs tries to predict the actions of the current student policy. This allows us approximate the actions of the student each expert state without additional simulation cost. The proxy student network is trained during the alignment phase with an LI loss on the network activations between proxy student and student, where both student and teacher observations are available for a subset of environment interactions.
Shared Action Decoder Rather than using separate networks to predict actions directly from obser- vations, we introduce shared action decoding layers used by the teacher, student, and proxy student. Each policy encodes its corresponding observations to features in common feature space, which are then fed to the action decoder to obtain the actions. Crucially, the action decoding layers are only updated based on the policy gradient computed with the task reward. The shared action decoder allows us to leverage the high-level feature correlations learned by the teacher through extensive environment interactions without requiring student observations. As a result, both the teacher and student learn to map their observations into common feature space instead of developing separate behavior policies. Furthermore, since the aligmment between teacher and student is trained using a limited subset of data, correctly simulating the true distribution of expert states be challenging.
4

---

Published as a conference paper at ICLR 2025
In this context, the shared action decoder can help it is trained broader set of expert states rather than the smaller alignment dataset, leading to more robust learning for the student.
4.2 TRAINING PHASES
Our proposed framework consists of three alternating training phases: (i) the classical policy roll- out, (ii) the policy update, and (iii) the alignment phase, see also Figure (b)-(c). The first two phases follow the standard on-policy training, while in the aligmment phase (ii)). the student F3s is aligned to the teacher Fr,, and the proxy student Fs aligned to the student Fs. For both network alignments, paired student and teacher observations are used. In the following, we provide more details about the specific training phases.
Roll-out Phase Our proposed framework introduces minor modification to the roll-out phase of the standard teacher training. specifically in the reward computation. addition to the task reward, we also add penalty term computed based the action difference between the teacher and the proxy student. This penalty encourages the teacher to only visit states in which the student can predict the same actions the teacher, thereby improving alignment between the student and teacher. Due to the addition of this term to the reward, it influences the long-term effect of taking an action on the divergence between student and teacher. Hence, it influences the exploration behavior of the teacher. In practice, during each roll-out phase, we store a subset of expert states required in the alignment phase. Depending on the simulation envrionment, this subset be randomly selected states fixed number of environments from which student observations are generated.
Policy Update Phase The gradient of the KL-Div term in Eq. 8 can be integrated into the policy up- date of the teacher, during which the network weights are updated using the clipped policy gradient of PPO. Since both the policy gradient and the KL--Diergence gradient are computed over the state distribution of the teacher, we can use the teacher states inside the roll-out buffer to compute the KL-Divergence between the action distributions of the teacher and the proxy student. This allows us to update the network parameters of the teacher in single backward pass through the combined loss function. Note that the KL--Divergence mentioned this phase is different to the KL-DDvergence penalty introduced at the rollout phase, and is directly integrated in the loss function to push the representation of the teacher to be similar to the student. In contrast, the penalty term affects the teacher's exploration by discouraging it from visiting regions of high disagreement with the student.
In the case of a continuous action space within PPO, we can simplify the KL-Divergence further to provide intuitive insights into the KL-Divergence gradient. In this setting, the teacher actions modeled as multivariate Gaussians of dimension d, with predicted mean jr, and state-independent covariances parameters Σr. Consequently, the KL-Divergence term in Eq. 8 simplifies to
E>
DKL (TT (8), (s(-)))))
= log
d+ Tr(E''
2
Ex
+ (() ps ps (s:)) (pr(s:) [s(ss))]-
The student and teacher share the same action decoder, which outputs mean action based on the latent state it receives. For the variance, we use state-independent parameter that is learned together with the action decoder based on teacher RL updates. Hence, the covariance matrix the same for the student and teacher, Σn = Es. As a result, the KL-Divergence reduces to the mean difference weighted by the covariance, along with an additional constant term
DKI(TT((1ss)) (=(+(s))) = cconst + (UT(s:) - us (s(())) ( - (s())))]
(10)
By optimizing Eq. 10, the teacher network Fr is aligned to the proxy student network Fg.. Intu- itively, the loss increases when the teacher confident in its actions (i.e. low Σr), but there still significant misalignment between the teacher and student. Consequently, the gradient update increases the covariance, leading to greater exploration during the roll-out phase, which increases the likelihood of the teacher discovering behaviors that are feasible for the student to learn.
Alignment Phase The alignment phase focuses on aligning the features across the encoders of the teacher, proxy student, and student. This phase is the only one that requires paired teacher and student observations, which are simulated from a subset of the teacher's s experiences during the roll- phase. We align the student encoder with the teacher encoder by computing the L.1 loss between
5

---

Published as a conference paper at ICLR 2025
Teacher without Alignment
Behavior Cloning
(c) DAgger
(4) Teacher with Aligmment
(c) Student without Alignment
(f) Student with Alignment
Figure 2: The goal of the agent is to navigate from the start (grey point) to the goal (green cell). The environment consists of four types of cells: empty (white), lava (red), and path (blue). The teacher can see all cell types while the student can not distinguish between lava and path. A teacher trained without alignment finds its optimal path through the maze (a), which can not be imitated by the student trained without alignment (e), Behavior Cloning (b), and DAgger (c). In contrast, teacher trained with alignment navigates around the maze (d), which can be easily copied by the student (f).
their corresponding features and between the activations of the frozen shared action decoder. To prevent the collapse of the model into predicting constant outputs, gradients are only backpropagated to the student encoder. Similarly, the proxy student is aligned with the student using the LI loss on the encoded features, with gradients only backpropagated the proxy student. The parameters of the teacher remain unaffected during this phase and are only updated during the policy update phase.
EXPERIMENTS
We compare method multiple baselines introduced below. We evaluate student-informed teacher training framework on three diverse tasks: maze navigation in a tabular setting, vision-based obstacle avoidance with quadrotor, and vision-based drawer opening using robot arm. Finally, we perform ablations to understand the role of different aspects of our method. For more details about the experimental setup and training, refer to the Appendix Section A.I.
Baselines We compare our method against multiple Imitation Learning baselines, which were trained with the same number of rendered images and environment interactions, including teacher training. Furthermore, all baselines use the same network architecture for the student and teacher.
Behavior Cloning (BC) The dataset, comprising privileged state information and rendered im- ages, is generated from the fixed number of timesteps of the teacher. The student network is then trained by minimizing the action loss between student and teacher on the collected data.
DAgger (Ross al., 2011) We employ an exponentially decreasing sampling threshold (8o = 0.98) to decide whether the actions of the teacher student are used at each timestep. Similarly to BC, we minimize the loss between the student and teacher actions on the collected samples.
HLRL (Radosavovic al., 2024) The weighting between the behaviour cloning loss and the vision-based RL loss is epxonentially decreased (5p 0.9998). To achieve successful runs, shared action decoder was necessary.
DHBC (Fu al., 2022) We treat the Privilged Info Encoder the teacher and the Adaptation Module as the student encoder, with the Unified Policy representing the shared action decoder. Both encoders are updated using their dual L1-loss in the feature space.
COSIL (Nguyen al., 2023) We adapt the method to PPO by using the L1 loss between student and teacher in the reward. As proposed in (Nguyen et al., 2023), the RL and IL objective is weighted with a learnable parameter, which is updated based on a target distance (0.5).
w/o Alignment (Ours) We exclude the KL-Divergence from both the reward and the policy update phase while keeping the paired L on shared action decoder features. As a result, the student does not affect the teacher.

---

Published as a conference paper at ICLR 2025
Aignment
Wthhou Aligmment
Posttion
v.d Algaraci
Start Povicun
Wiee Alganei
3.0
0.0
-15
-10
[m]
Figure 3: Vision-Based Obstacle Avoidance. On the left, the teacher trajectory rollouts of the vision-based quadrotor obstacle avoidance tasks are visualized. Our approach results in policy behavior where the quadrotor adjusts the camera's 's viewing direction to capture sufficient environ- mental information for the student policy.
Implementation Details For the vision-based tasks, we use a three-layer MLP with ELU activations for the teacher and the proxy student, which both receive privileged observations. The shared action decoder consists of a single fully connected layer. The student processes images using frozen DINOv2 encoder (Oquab al., 2023). These image features are combined with state observations and passed through an additional MLP before being forwarded to the shared action decoder.
5.1 ENVIRONMENTS
Color Maze We first test method maze navigation task in tabular setting. The objective of the agent is to reach the goal (green cell) from the starting point (grey point), as visualized in Figure 2. In the center of the environment
maze where the agent can only navigate along randomly generated path (blue cells), and the episode ends if the agent steps into the lava (red cells) or reaches the target. The teacher observations include the distance the goal, the type of neighboring cells (empty, lava, or path), and the relative movement from the previous timestep. The student receives the same observations, except that the cell types are classified as empty or occupied, which maps lava and path cells to the same value, which leads to teacher-student asymmetry. The reward structure assigns a positive reward for reaching the target (10) and for visiting each new path cell (0.5), while penalties are given for moving into lava (0.1) and revisiting path cells (0.5).
Vision-based Obstacle Avoidance with a Quadrotor We evaluate our approach on the complex task of agile quadrotor control, in which quadrotor needs to follow constant velocity command while avoiding obstacles. To simulate this, we design a custom environment using the realistic agile quadrotor simulator, Flightmare (Song al., 2021). The environment is world box measuring 30m x 30m x 3m, within which we randomly place four static pillar obstacles, each with a radius of 1.5m. At the start of each episode, the quadrotor's position and viewing direction are randomly initialized at collision-free locations. Next, we sample a velocity command for the quadrotor that would result in potential collision with one of the obstacles. The teacher policy receives observa- tions composed of the quadrotor's 's linear velocity, commanded linear velocity, orientation, angular velocity, and the relative distance to each obstacle. In contrast, the dision-based student policy re- places obstacle positions with inputs from an RGB camera with a limited field of view, as shown in Figure 3. The reward is detailed appendix 1.2. During our experiments, we conduct 256 runs with random initialization, each running for 1000 steps for all the baseline approaches.
Vision-based Manipulation We further evaluate our method on a complex vision-based manipula- tion scenario, where a robot arm learns to open drawer. We adapt the publicly available Omniverse Isaac Gym Reinforcement Learning Environments for Isaac Sim repository, specifically modifying the cabinet-opening task. In this task, a Franka robot arm is trained to open drawer that contains objects inside. The teacher receives observations composed of the robot's 's joint positions and veloc- ities and the relative position between the gripper and the drawer handle. The student observations
7

---

Published as a conference paper at ICLR 2025
Methods
Success Rate
Stidest Retun
Tescher Return
Algesce
Algence
Widac Algunoa
BC
0.05 + 0.04
Avenze Eptsode Retim
50
DAgger
0.08 = 0.03
HLRL
0.31 = 0.11
DWBC
0.35 + 0.07
COSIL
0.30 + 0.07
30
20
w/o Align (Ours)
0,38 + 0.11
10:00
3000
1040
4000
5000
0400
2030
1040
4.00
10.0
Align (Ours)
0.46 + 0.04
Fiechs
Fjech
Obstacle Avoidance Suc- cess Rates. The mean and stan- dard deviation of the success rate for vision-based quadrotor flight ob- tained from three trainings.
Table I:
Figure 4: Obstacle Avoidance Returns. The mean returns achieved by the student and teacher trained with and without alignment, averaged over three training runs.
replace relative distance information with an image. Instead of an unobstructed top view, the view- point of the camera is selected closer to the robot arm. This camera setup closer to real-world platforms, such as humanoid robots, where certain arm configurations may obstruct cameras.
5.2 RESULTS
Color Maze In Figure we show single trajectory (green path) and the occupancy grid (blue cells) for 15 evaluation environments across five training runs for each tested policy. As can be seen in Figure 2, students trained with Behavior Cloning (BC) or DAgger fail to imitate the behavior of the teacher due to their inability to differentiate between lava and path cells. In Figure 2 (c), we observe a similar behavior for the student trained with shared action decoder but without teacher alignment. In contrast, when using our framework, with a teacher penalty for visiting states where the student is unable to predict the same action, i.e., inside the maze, the teacher learns to avoid the maze and navigates around to reach the target. This behavior is successfully imitated by the student, who reaches the target in all test runs, unlike all other baselines. These results confirm that framework helps the teacher learn strategies that easier for the student imitate, such as avoiding the maze altogether. Additionally, since the behavior of navigating around the maze leads to sparse rewards once the goal is reached, the results demonstrate that the penalty term does not hinder the exploration necessary to discover the optimal student solution.
Vision-based Obstacle Avoidance with Quadrotor The results shown in Table 1 clearly show that our approach with alignment achieves the highest success rate with 0.46 compared to 0.08 (DAgger) and 0.05 (BC). Moreover, we show that our approach surpasses another baseline ap- proach DWBC by 0.11, 1, HIRL by 0.15, and COSIL by 0.16, emphasizing that our adaptive teacher training offers more effective guidance to the student than the integration of RL rewards and with a static teacher. Table 1.1.3 highlights the perception awareness metrics, showing our approach im- proves the student policy's behavior and performance. This is due to the fact that, the teacher policy lacks "perception awareness of obstacles," meaning the obstacles are not fully observable during the avoidance task from the equipped camera. As result, the vision-based student struggles to infer the correct actions, leading to very low success rates (below 0.1), despite all teacher policies achieving over 98% success rates. This highlights the importance of enforcing information symmetry between teacher and student policies. The reason the success rates for the imitation baselines are not zero is that, in some configurations, the quadrotor is initialized with a viewing angle that allows the student to detect obstacles, enabling it to replicate the teacher's policy. Figure illustrates the rollout trajec- tories, where we observe that the policy learned through our approach adjusts the viewing direction during flight align with the velocity direction. This ensures that all encountered obstacles are vis- ible in the camera's 's field of view, allowing the student to infer the necessary actions using sufficient information. Figure 4 shows the returns of our method with and without alignment for both the teacher and student policies. The results show that our proposed alignment is crucial for the student to reach reward regions similar the teacher's 's despite the student's more limited observability.
basedd Manipulation Table 3 reports the success rates of the student for our method (with and without alignment) and the baselines across five training runs. Our framework with and with- out alignment significantly outperforms students trained with DAgger and BC, with success rates of up to 0.77 compared to 0.47 (DAgger) and 0.27 (BC). This improvement be explained by the
8

---

Published as a conference paper at ICLR 2025
Task Environment
Without A Alignment
With Alignment
Figure 5: Vision- Based Manipulation. On the left, the task of opening a drawer with a robotic arm is visualized for all of the parallel environments. The two images in the center (Without Alignment) are sample images given to the student, which show the teacher behaviors trained without our align- ment. Our approach with alignment leads behaviors that the student can imitate more easily, i.c., the robot does not block the red drawer handle, as visualized in the two images on the right.
shared task decoder, which is trained by leveraging the training samples of the teacher. Additionally, our method outperforms DWBC, HLRL, and COSIL with a success rate improvement in the range of 0.25-0.32, which confirms that the adaptation of the teacher provides better student supervision than the combination between RL reward and IL with a fixed teacher. Our method with alignment improves student success rates by 17% compared to the non-aligned framework while also consis- tently achieving higher returns (Figure 6). These results demonstrate that our framework helps the teacher Iearn better behaviors for the student. Student policies trained without alignment achieve non-zero success rates due to the small sampling interval of the cabinet position. This allows them to memorize behaviors without relying heavily on the images. All teachers, regardless of alignment, reach 100% success rate. Interestingly, the return of the teacher trained with alignment is also constantly higher than without alignment. A possible explanation is that the teacher learns gripper movements optimized for robustness without trying multiple times to grab the handle, which is difficult behavior for the student lacking relative pose information.
As be seen in Figure 5, our method leads to several different behaviors. With alignment, the teacher learns once to grab the handle from a top-down configuration while another teacher lowers its first two links to make the red handle visible. In both cases, the red handle is visible right before the gripper touches it. This shows that our alignment leads to emerging teacher behaviors that consider the imitation difficulties of the student.
Ablations To validate our design choices, we perform multiple ablations on the manipulation envi- ronment, reported Table 2. Namely, we study the effect of the penalty term (w Penalty), the KL- Divergence gradient in the policy update (w KL-Grrad), and the role of the shared action decoder. Our results indicate that the KL.-Divergence gradient and the shared action decoder are crucial to the large improvement achieved by method. As for the alignment penalty, its absence leads to a 14% drop in performance. Finally, we study the sensitivity to the coefficient A of the alignment penalty. Our results show clear robustness to this choice, as method consistently achieves high success rate regardless of the choice. Setting A1 = 0.025 leads to an impressive 95% success rate.
6 RELATED WORK
Imitation learning from teacher policy has shown significant progress across various domains. Approaches such as behavior cloning (BC) (Bain & Sammut, 1995) and DAgger (Ross et al., 2011) typically train on a dataset of encountered states paired with expert actions. The core idea is train a student policy to imitate the behavior of either a reinforcement learning (RL) policy that has been trained using state information objective is to learn policy that closely matches expert demonstrations through supervised learning. DAgger (Ross al., 201 1) improves upon this by aggregating expert demonstrations with the student policy's experiences during training to enhance generalization beyond the original demonstrations. Ross & Bagnell (2014) extended DAgger by taking into account not only the expert's actions but also the long-term consequences of the student's 's policy quantified using the task cost. Ho & Ermon (2016) propose a method for matching the state-action distribution using a generative adversarial algorithm similar to the one used in generative adversarial networks (Goodfellow al., 2014). Other methods explored leveraging priviliged information in components that are only relevant for training such crities, reward models and dynamics models (Pinto et al., 2018; 8; Hu et al., 2024).
through expert demonstrations. In BC (Bain & Sammut, 1995), the
One key assumption shared by these methods is that the observation information available to the student policy must be at least comprehensive that available the teacher (Osa et al., 2018;

---

Published as a conference paper at ICLR 2025
w Penalty/ / w/o KL- Grad
wlo Penalty/ /
wlo shared decoder
(Default)
w KL-Grad
A 0.1
A = 0.075
A = 0.05
A = 0.025
Success Rate
0.47 + 0.37
0.74 + 0.08
0.62 0:27
0.81 0.20
0.77 + 0.18
+: 0.07
0.95
Table 2: Ablations for Vision- Based Manipulation. The mean and standard deviation of the suc- cess rate obtained from five trainings runs.
Methods
Success Rate
Student Retum
Teacher Return
1750
BC
0.16 = 0.15
1500
DAgger
0.34 + 0.31
HLRL
0.61 + 0.22
1250
DWBC
0.63 + 0.18
1009
COSIL
0.56 + 0.21
With Alignntent
w/o Align (Ours)
0.61 + 0.18
Without Aliggment
Align (Ours)
0.88 + 0.07
500
Fpochs
Flwccs:
Manipulation Success Rates. The mean and standard deviation of the success rate for the task of opening the drawer obtained from five trainings runs.
Table 3:
Manipulation Returns. Mean returns of the student and teacher trained with, without alignment and DWBC, averaged over five training runs.
Figure 6:
Swamy et al., 2022). This assumptions falls short in many scenarios, especially when the student operates in a partially observable setting. In such cases, the student could easily struggle to infer the teacher's behavior based on its own partial observability, Hence, multiple methods proposed aug- ment the teacher distillation objective with a student RL objective (Weihs et al., 2021; Nguyen et al., 2023; Shenfeld et al., 2023; Radosavovic et al., 2024). These methods introduce different ways to balance the imitation and RL terms, either using fixed trade-off coefficients (Nguyen et al., 2023), fixed schedule (Radosavovic al., 2024), or some elaborate heuristic (Weihs et al., 2021: Shenfeld et al., 2023). Similarly, Walsman et al. (2022) jointly train two separate policies besides the teacher: a follower, which learns to follow the teacher, and an explorer, which maximizes environmental re- wards using the follower's value function for reward shaping. Alternatively, in an off-policy setting, the agent can directly optimize the student policy with RL using interactions collected by rolling out the teacher to facilitate exploration (Chane-Sane et al., 2024). In contrast to these methods which focus on guiding the student exploration, we approach the challenge of asymmetric learning by regu- larizing the teacher policy. A similar concept have been explored for multi aggti- self-play (Hamade et al., 2024) and student-teacher training (Warrington et al., 2021; Fu et al., 2022). Namely, Warring- et al. (2021) propose exposing the teacher training to student rollouts ground the behaviors learnt by the teacher. Fu et al. (2022); He et al. (2024) learn a unified student-teacher policy and two separate encoders which are encouraged to extract similar features. In the context of generative models, Bachman & Precup (2015) proposes to jointly optimize a primary policy and a guide policy to reconstruct image regions. Our approach augments the teacher loss by adding a loss that penalizes confident actions in areas where the teacher-student discrepancy is high. Additionally, we modify the teacher's 's reward to discourage visiting states where the student struggles to infer teacher actions.
7 CONCLUSION
This paper tackles the problem of teacher-student asymmetry in imitation learning. To address this problem, we propose novel method that trains the teacher not only based on task objectives but also by encouraging behaviors that the student can successfully imitate. By including the performance bound in imitation learning in the teacher objective, we derive two key modifications to the teacher training: a reward term penalizing states with a high teacher-student misalignment and directly aligning the teacher the student based on the KL-Divergence between their action distributions. Our approach is not limited to specific IL algorithms and can seamlessly integrate with various II. pipelines, including multi-agent II. We illustrate our method a tabular maze task, where outperforms standard imitation learning baselines that fail to reach the target. Furthermore, our framework enhances student performance in the complex task vision-based quadrotor obstacle avoidance, where perception-aware flight emerges without the need for explicit reward tuning. We also validate our method in a vision-based manipulation task, showing how the teacher learns to open drawer without obstructing critical visual information in the camera view of the student. We believe that our framework extends privileged imitation learning to various modalities with large information gaps, unlocking the potential to tackle complex real-world tasks in robotics.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Elie Aljalbout, Felix Frank, Maximilian Karl, and Patrick van der Smagt. On the role of the action space in robot manipulation learning and sim-to-real transfer. IEEE Robotics and Automation Letters, 2024.
Philip Bachman and Doina Precup. Data generation as sequential decision making. Advances in Neural Information Processing Systems, 28, 2015.
Michael Bain and Claude Sammut. A framework for behavioural cloning. In Machine Intelligence 15, pp. 103-129, 1995.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
Elliot Chane-Sane, Joseph Amigo, Thomas Flayols, Ludovic Righetti, and Nicolas Mansard. Soloparkour: Constrained reinforcement learning for visual locomotion from privileged expe- rience. In Conference on Robot Learning, 2024.
Dian Chen, Brady Zhou, Vladlen Koltun, and Philipp Krähenbühl. Learning by cheating. In Con- ference on Robot Learning (CoRL), 2019.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. Robotics: Science and Systems, 2023.
Ria Doshi, Homer Walke, Oier Mees, Sudeep Dasari, and Sergey Levine. Scaling cross-embodied learning: One policy for manipulation, navigation, locomotion and aviation. 8th Conference on Robot Learning, 2024.
Zipeng Fu, Xuxin Cheng, and Deepak Pathak. Deep whole-body control: Learning unified policy for manipulation and locomotion. In Conference on Robot Learning (CoRL., 2022.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Karim Hamade, Reid Mcllroy-Younr Young, Siddhartha Sen, Jon Kleinberg, and Ashton Anderson. De- signing skill-compatible AI: Methodologies and frameworks chess. In The Twelfth Interna- tional Conference Learning Representations, 2024. URL htctps://openrevieee net/
Forum?idee serum?id=79rfqeee 1w4.
Haoran He, Peilin Wu, Chenjia Bai, Hang Lai, Lingxiao Wang, Ling Pan, Xiaolin Hu, and Weinan Zhang. Bridging the sim-to-real gap from the information bottleneck perspective. In 8th An- nual Conference on Robot Learning, 2024. URL http::/:: /tt. Bq4X0aU4sV.
cont/forummid
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances neural information processing systems, 29, 2016.
Edward S. Hu, James Springer, Oleh Rybkin, and Dinesh Jayaraman. Privileged sensing scaffolds reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024. URL https: / http///ppppteeeeee/e/eeeeeeeeeet
Elia Kaufmann, Leonard Bauersfeld, Antonio Loquercio, Matthias Müller, Vladlen Koltun, and Davide Scaramuzza. Champion-level drone racing using deep reinforcement learning. Nature, 620(7976):982 987229977111 2023.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations (ICLR), San Diego, CA, USA, 2015.
Jianlan Luo, Charles Xu, Xinyang Geng, Gilbert Feng, Kuan Fang, Liam Tan, Stefan Schaal, and Sergey Levine. Multi-stage cable routing through hierarchical imitation learning. IEEE Transac- fions Robotics, 2024.
11

---

Published as a conference paper at ICLR 2025
Denys Makoviichuk and Viktor Makoviychuk. rl-games: A high- perf-peeformance framework for rein- forcement learning. https:
www.///nwwnn////eeeee games, May 2021.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 (18(7540):529- 533,, 2015.
Hai Huu Nguyen, Andrea Baisero, Dian Wang, Christopher Amato, and Robert Platt. Leveraging fully observable policies for learning under partial observability. In Conference Robot Learn- ing. pp. 1673-1683. PMLR, 2023.
Maxime Oquab, Timothée Darcet, Theo Moutakanni, Huy V. Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fermandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, Russell Howes, Po-Yao Huang, Hu Xu, Vasu Sharma, Shang-Wor Wen Li, Wojciech Galuba, Mike Rabbat, Mido Assran, Nicolas Ballas, Gabriel Synnaeve. Ishan Misra, Herve Jegou, Julien Mairal, Patrick Labatut, Ar- mand Joulin, and Piotr Bojanowski. Dinov2: Learning robust visual features without supervision, 2023.
Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J Andrew Bagnell, Pieter Abbeel, Jan Peters, al. An algorithmic perspective on imitation learning. Foundations and Trends@ in Robotics, 7(1-2): 1-179, 2018.
Kei Ota. Tomoaki Oiki, Devesh Jha, Toshisada Mariyama, and Daniel Nikovski. Can increasing put dimensionality improve deep reinforcement learning? In International conference on machine learning. pp. 7424-7433. PMLR, 2020.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch, 2017.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning, pp. PMLR, 2017.
Lerrel Pinto, Marcin Andrychowicz, Peter Welinder, Wojeiech Zaremba, and Pieter Abbeel. Asym- metric actor critic for image-based robot learning. Robotics: Science and Systems, 2018.
Ilija Radosavovic, Tete Xiao, Bike Zhang, Trevor Darrell, Jitendra Malik, and Koushil Sreenath. Real-world humanoid locomotion with reinforcement learning. Science Robotics, 9(8))ea 2024.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021. URL http://jmlr.... org/ pttp://appsss/221 20-12 10--1364.20 html.
Stephane Ross and Andrew Bagnell. Reinforcement and imitation learning via interactive no-regret learning. arXiv preprint arXiv:1406.599 2014.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and struc- tured prediction to no-regret online learning. In Proceedings of the fourteenth international con- ference artificial intelligence and statistics, pp. 627 635. JMLR Workshop and Conference Proceedings, 2011.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1.X7:0777.00347,, 2017.
Idan Shenfeld. Zhang- Wei Hong, Aviv Tamar, and Pulkit Agrawal. Tgrl: An algorithm for teacher guided reinforcement learning. In International Conference on Machine Learning, pp. 31077- 31093. PMLR, 2023.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7787):484 489, 2016.
12

---

Published as a conference paper at ICLR 2025
Yunlong Song, Selim Naji, Elia Kaufmann, Antonio Loquercio, and Davide Scaramuzza. Flight- mare: A flexible quadrotor simulator. In Proceedings of the 2020 Conference on Robot Learning, pp. 1147-11577 2021.
Gokul Swamy, Sanjiban Choudhury, J Bagnell, and Steven Z Wu. Sequence model imitation learn- ing with unobserved contexts. Advances in Neural Information Processing Systems, 35:1 35:17665- 17676, 2022.
Umar Syed and Robert E Schapire. A reduction from apprenticeship learning to classification. Advances in neural information processing systems, 23, 2010
Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et al. Octo: An open-source generalist robot policy. Robotics: Science and Systems, 2024.
Oriol Vinyals, Igor Babuschkin, Wojeiech Czarnecki, Michaël Mathieu, Andrew Dudzik, Juny- oung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. nature, 575(7782):350-3344 2019.
Aaron Walsman, Muru Zhang, Sanjiban Choudhury, Dieter Fox, and Ali Farhadi. Impossibly good experts and how to follow them. In The Eleventh International Conference Learning Repre- sentations, 2022.
Andrew Warrington, Jonathan W Lavington, Adam Scibior, Mark Schmidt, and Frank Wood. Robust asymmetric learning in pomdps. In International Conference Machine Learning, pp. 11013- 11023. PMLR, 2021.
Luca Weihs, Unnat Jain, lou-Jen Liu, Jordi Salvador, Svetlana Lazebnik, Aniruddha Kembhavi, and Alex Schwing. Bridging the imitation gap by adaptive insubordination. Advances in Neural Information Processing Systems, 34:19134-19146, 2021.
Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 15737-15749. Curran Associates, Inc., neurips .
2020. URL htttpps:/pproceeednnge htte://5501100011100066666060233
A APPENDIX
A.1 TRAINING DETAILS
A.1.1 COLOR MAZE
We implemented the color maze environment using the Gym framework (Brockman et al., 2016), with the discrete agent trained via the PPO (Schulman et al., 2017) implementation of StableBase- lines3 (Raffin et al., 2021) based on Pytorch (Paszke et al., 2017). The training setup involves 1,000 parallel environments, each containing randomly generated path through the maze located the center of 21x21 grid. The maze itself spans 15x15 cells, leaving a three-cell-wide empty border around through which the optimal path of the student goes. The paths are generated randomly by alternating between vertical segments of 3 to 5 cells and horizontal segments of 1 to 4 cells. Given the large number of environments, each maze path is generated only once at the start of training and remains fixed throughout the process to save computation time.
The action space of the agent consists of four discrete movements: one cell up, right, down, or left. If the agent selects action that would move it outside the grid, its position remains unchanged. Moving into a lava cell results in immediate episode termination and a negative reward. In contrast, reaching the target cell terminates the episode with positive terminal reward.
The teacher training uses in total 10a environment interactions. During each roll-out phase, 250 steps collected from each of the 1,000 parallel environments, resulting in 250,000 experiences. In the policy update phase. the teacher network is updated using a batch containing all 250,000 col- lected experiences. In addition the default PPO gradients, apply the KI KL--Divrgencc gradient
13

---

Published as a conference paper at ICLR 2025
with weighting coefficient of 0.001. The entropy coefficient for PPO is set to 0.3. During each paired alignment step, the observations stored the roll-out buffer are used as a single batch to update the imitated student and student networks over 20 iterations. The networks using the L1 loss between corresponding features, with network weights updated via the ADAM optimizer (Kingma & Ba, 2015).
optimized
The reward function consists terminal rewards for either reaching the target (fsucoess) or stepping into a lava cell (ffa)).. To encourage the teacher to follow the correct path through the maze, we assign a reward when the agent moves to path cell for the first time pa To prevent the teacher from simply cycling over already-
path cells, we impose a penalty for revisiting the same path cell reevisit path- In the case of our alignment framework, an additional penalty term based on the Kl.-Divergence between teacher and student actions included the reward function L-L.
Tto 10r success + 0.5rpath revisttssit.p falr 1.9TKL.
(11)
A.1.2 VISION-BASHI QUADROTOR FLIGHT
For the obstacle avoidance task, we build a customized RL training environment using Flightmare and Stable-baselines..
this task, we define the observations as Oobstacle R. Vemd: V, w. aprev 6p1, 8p2, 5ps, 6ps.. where R R is vector comprising the first two columns of R, E IR and represents the previous action from the actor policy, and ópp I22 represent the relative dif- ference in position of Here, ôp; represents the absolute distances of the quadrotor to all obstacles represented in the world frame. The total reward time t,, denoted as rt, consists of several compo- nents:
R$ denote the linear and angular velocity of the drone, Oprev
Tobstacle Uprog + Fact + Fbr + Tobstacle crash + Terash
(12)
where Tprog represents progress toward command velocity direction to ensure velocity following, Fact penalizes changes in actions from the previous time step, Thr discourages high body rates to ensure stable flying behavior, Ppsss is a penalty for crashing encountering obstacles, weighted by the velocity. and Ferash is a crashing penalty when it collides to the ground or flies out of the world box. The policy trained using standard PPO implemented in the stable- baseliess library. The reward components are formulated as follows
r Pprog 2.0 Σ Ep.
Pact 0.15/la, a--1l
(13)
Tbr = 0.05||wBBell
Fobstacle crash 10 if robot erash to the obstacle,
I Terash = 5|||i||
if robot crashes (ceiling, ground).
Over the course of training, the teacher collects data across 6300 epochs, with 100 parallel environ- ments. Among these environments, 64 render images at each timestep. During the policy update phase, we use a minibatch size of 12500.
A.1.3 VISION-BASED MANIPULATION
For the manipulation task, the Omniverse Isaac Gym Environments framework. Specifically, we modify the existing Franka Cabinet Task, where a Franka arm is tasked with opening the top drawer of cabinet. We make several minor modifications color of the cabinet to brown and the top drawer handle to red to improve visibility and highlight when the handle is occluded. Second, we remove the knobs from the bottom doors and the handle from the bottom drawer to eliminate alternative visual cues for locating the top handle. To increase occlusion, we raise the Franka arm by 0.4m and rotate the cabinet (60 degrees around the z-axis), as shown in Fig. 5. Additionally, we uniformly sample the x and y positions of the cabinet relative to the robot arm within a grid of [-0.15m, 0. 0.15m] x [-0.Im, 0.1 The initial position of the robot
this setup. First, we change the
chosen to ensure that the default teacher consistently occludes the camera while still maintaining enough distance between the robot arm and the drawer handle to ensure a certain degree of exploration.
14

---

Published as a conference paper at ICLR 2025
The teacher agent is trained using the RL-Games framework (Makoviichuk & Makoviychuk, 2021), specifically with a PPO continuous agent. Over the course of training, the teacher collects data across 1,500 epochs, with horizon of 16 timesteps and 4,096 parallel environments, resulting in total of 98,304,000 environment interactions. Among these environments, 128 render images at each timestep. generating 3,072,000 data samples for the paired alignment phase. To manage these image samples, we employ a rolling buffer with a size of 100,000, updating it in a FIFO (First In, First Out) manner. During the policy update phase, we use minibatch size of 8,192. The KL-Divergence loss is weighted by 0.01 in the policy update. For all other PPO hyperparameters, we use the default settings provided by the Omniverse Isaac Gym framework.
The action space of the robot arm is the multi-step integration joint position Aljalbout et al. (2024). The Panda arm has joints and the gripper has fingers, making for total of 9 degrees of freedom. For the teacher, we use the default observation set, which includes the joint positions, joint velocities, the relative 3D distance between the gripper and the drawer handle, as well the 3D position and velocity of the drawer joint. In contrast, the student receives only the joint positions and angular velocities, along with an additional image input, as shown Fig. 5, We use the default reward formulation, which includes components such as the distance between the gripper and the handle, the gripper's s orientation, and the position of the gripper's "thumbs," among others. For our alignment training, the KI.-Divergence is weighted by 0.05 and added to the task reward.
joint
A.2 ALGORITHM OVERVIEW
A pseudocode description of the different training phases with the crucial steps is provided in Algo- rithm!
Algorithm 1: Student-Informed Teacher Training
Input:
Teacher policy Tr, Student policy æs, Proxy student policy øs. Task environment E, Reward function KL-divergence weight Y₁ A2, Iterations N, Roll-out steps T', Policy update steps M, Alignment steps L
Output: Trained teacher Ty and student IT Initialize: Teacher, student, proxy encoders, and shared task decoder
Buffers: Experience buffer Beep. Alignment buffer Batign
for to
// Roll-Cut Phase
for t + 1o T do
Sample action a; TT((s)
Execute a, to get rt, 8141
Update reward: r Ti (T
Add (St, ae, rh, 8t41) to Bexp
Render Osud for subset of visited states and add Balign
end
// Policy Update Phase for j 1 to M do
Sample experiences from Beep Compute loss: L Epolicy DRDDII(T(((ss) Update teacher policy गT using L
end // Alignment. Phase for k - 1 to do
Sample (0T,0s) Balign
Update πg by aligning teacher and student features
Update ãs by aligning student and proxy features
end end
15

---

Published as a conference paper at ICLR 2025
A.3 EVALUATION ON LEARNING PERCEPTION AWARENESS
To further demonstrate the effectiveness of our approach, we evaluated how the resulting policies essentially perception-aware. In experiments, we argue that the vision-based student policy in our framework can learn to fly in a manner that effectively perceives the environment using the onboard camera. To validate this, used two different metries: the Velocity Angle and Number of Obstacles in View. The Velocity Angle measures the angle between the quadrotor's 's heading (camera viewing direction) and its velocity direction. Ideally, the quadrotor should "look at" the direction it is flying towards. A smaller angle ensures that obstacles are perceived in time, allowing the controller to gather sufficient information to act effectively. Additionally, we introduced another direct metric, the Number of Obstacles in View, which quantifies how many obstacles the policy can perceive during each rollout. As shown in A.1.3, our approach demonstrates superior performance in terms of perception awareness compared to all baseline methods. This further underscores the effectiveness of our approach.
Velocity Angle [deg] 4
Num. Obstacle in
Methods
View f
BC
75.5
1.92
DAgger
78.6
2.42
HLRL
61.9
2.09
DWBC
46.7
2.33
COSIL
69.3
2.24
w/o Align (Ours)
63.2
2.61
Align (Ours)
32.2
3.51
16


Reviews:
[{"review_id": "eb708io9S2", "reviewer": "Reviewer_2qRa", "summary": "The authors propose to improve imitation learning between a privileged expert policy that sees more information and a partially observed student policy. They propose to train the teacher policy so that it maximizes return while staying close to the student policy distribution, so that the teacher and student trajectory distributions are more aligned. The authors show this improves performance in several simulated robotic setups.\n\n---\nPost rebuttal - the authors have done a decent job addressing my concerns on experimentation and improved the presentation.\n\nUpdated my score to 6.\n\nI do agree with reviewer mh2n that this submission still feels a bit rushed, with a lot of important results (comparison to competitive baselines and ablations) coming in through the rebuttal period. I tend to be a bit more wary of results that come in through rebuttal, given the limited time and resources.", "questions": "See weaknesses above. I am eager to see the paper be improved.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- Simple approach to handling asymmetry in teacher student policy distillation, where teacher is trained to minimize divergence between teacher and student. this will constrain the teacher to states where the student can explore\n- Interesting experimental findings.\n\t- The experiments show that teacher policies trained in this manner are better for teaching student policies.\n\t- A teacher with this constraint gets higher returns than a teacher without, because learning to act with less inputs leads to more robustness and generalization behavior. This would be really interesting to explore further, with more experiments or analysis.", "weaknesses": "## Experimental section has several problems\n- Experimental section is lacking, sparse in task selection and choices of baseline, and seem a bit contrived. Several ways to fix this:\n\t- More standardized benchmarks, using envs from prior work like  [1, 2]\n\t- better baselines and analysis (see below)\n\t-  Sim2real of the drone / manipulator results would round it out\n\n- choice of baselines is lacking, really should compare with prior work in handling asymmetric RL / IL problems ( see references below)\n- Baseline details are completely missing, even if BC and Dagger are widely known, their training details need to be written down.\n- little-to-no ablations, i.e. why the shared feature encoder? Wouldn't we expect that the teacher representation be very different than the student representation in very partially observed settings? \n\n- The authors need to put more work in the experimental section, which is verbose and unorganized. Having two giant paragraphs per experiment, one for setup, and one for results, is lazy and hard for readers to parse. Would like to see more figures, analysis of the teacher and student behavior, especially in the drone and manipulator case. IMO, the color maze experiment, which is a toy experiment, does not require that much text and space dedicated to explaining it. It could even be moved to appendix.\n- Another way to improve the experimental section, is to provide some real world robot results. Because this is ICLR (more focused on learning methods), I would say this isn't a strict requirement, and I would appreciate more rigorous comparisons in the previous experiments. But real robot results would definitely round out the experimental section.\n\n## Another weakness - having to train a teacher policy from scratch\n- One drawback of this paper is the need to train a teacher policy with their specialized objective. This makes the method harder to use in practice since this method requires training a teacher policy from scratch using RL and many, many samples in a fast simulator  (10^8 for maze, 100M for manipulation). Getting a good simulator and doing sim2real is not always feasible.\n\t- In contrast, other methods like ADVISOR, COSIL, etc. (see references below) assume that the teacher policy is given, and do not require teacher training. In robotics,  where reasonable teacher policies can be obtained through scripting and pre-existing controllers,  it seems that approaches that do not modify the teacher are easier to use. \n\nSome discussion here would be appreciated.\n\n\n## Minor: missing connections to prior work on regularizing privileged teachers and representations\n- Regularizing the privileged policy / representation so that it doesn't stray too far away from the student  has been explored in the past, see [6,7,8]\n\n1. Leveraging Fully Observable Policies for Learning under Partial Observability\n2. Privileged Sensing Scaffolds Reinforcement Learning\n3. Bridging the Imitation Gap by Adaptive Insubordination\n4. TGRL: An Algorithm for Teacher Guided Reinforcement Learning\n5. Impossibly Good Experts and How to Follow Them \n6. Bridging the Sim-to-Real Gap from the Information Bottleneck Perspective\n7. Deep Whole-Body Control: Learning a Unified Policy for Manipulation and Locomotion \n8. Designing Skill-Compatible AI: Methodologies and Frameworks in Chess", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "14HSYAkgdw", "reviewer": "Reviewer_xapm", "summary": "The paper addresses the important problem of information asymmetry between teacher and student policies in imitation learning. The authors propose a novel joint training framework that encourages the teacher to learn behaviors that can be more easily imitated by the student, thereby mitigating the negative impact of the information gap. The core idea is to incorporate the performance bound of imitation learning into the teacher's objective function, resulting in two key modifications: a KL-divergence-based penalty term in the teacher's reward function and a KL-divergence-based supervisory signal for updating the teacher's network parameters. The effectiveness of the proposed method is validated through experiments on three diverse tasks: maze navigation, vision-based quadrotor obstacle avoidance, and vision-based drawer opening with a robotic arm.", "questions": "1. How does the proposed method theoretically compare to existing approaches in terms of convergence guarantees and sample complexity?\n2. What is the sensitivity of the method to the choice of the KL-divergence penalty weight, and how can this hyperparameter be effectively tuned?\n3. Can the authors provide more detailed ablation studies to demonstrate the individual contributions of the reward penalty and the KL-divergence supervision?\n4. How well does the method scale to more complex environments and tasks, and what are the potential limitations in terms of computational overhead and sample efficiency?\n5. Can the authors discuss the potential extensions of the proposed framework to other imitation learning algorithms and settings, such as inverse reinforcement learning or multi-agent imitation learning?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents a novel approach to tackling the teacher-student information asymmetry problem by incorporating the imitation learning performance bound into the teacher's objective function, leading to a creative combination of ideas from imitation learning theory and practical algorithm design.\n2. The proposed method is well-motivated and grounded in the theoretical foundations of imitation learning, with clear explanations for the design of the KL-divergence-based penalty and supervisory terms.\n3. The method is evaluated on three diverse tasks, demonstrating its applicability to both discrete and continuous control domains. The results show improvements over baseline imitation learning approaches.\n4. The paper is generally well-written and easy to follow, with a clear problem statement, detailed method description, and effective use of figures and tables.", "weaknesses": "1. Limited theoretical analysis: The paper lacks a comprehensive theoretical analysis comparing the proposed method to existing approaches. A more rigorous theoretical justification for the superiority of the method would strengthen the contributions.\n2. Insufficient ablation studies: The individual contributions of the reward penalty and the KL-divergence supervision are not clearly distinguished through ablation experiments, making it difficult to assess the necessity of each component.\n3. Limited experimental evaluation: While the method is evaluated on three tasks, a more extensive experimental evaluation on a wider range of environments and benchmarks would provide stronger evidence for the generalizability and effectiveness of the approach.\n4. Clarity of certain technical details: Some aspects of the paper, such as the definition of the proxy student network and the justification for the equality of covariance matrices in equations (9) and (10), require further clarification.", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "e6YXMFOBuk", "reviewer": "Reviewer_mh2n", "summary": "The paper proposes to tackle the problem of information asymmetry in the privileged training and distillation paradigm by proposing to adapt the teacher policy’s behavior to account for the student’s partial observability by adding a reward term to encourage imitability of actions. They additionally introduce a proxy student network that approximates how the student behaves conditioned on the teacher’s privileged observations to alleviate the need of generating potentially high-dimensional student observations for optimizing the teacher policy.", "questions": "(Q1) The description of the alignment phase (Sec 4.2) can benefit from more clarity: What subset of experiences are good enough to train the proxy student – how much divergence between student/student-proxy is tolerable? If this has to be a large subset of the data then there is no benefit from using a proxy student, one might as well use all the observations to update both the teacher and student. I think clarity of the approach will improve with a pseudocode description of all the phases with input requirements in the appendix.\n\n(Q2) In the experiments, the baselines aren’t described clearly – what does w/o alignment mean? Does it suggest removal of proxy student network in (a) no imitability loss for teacher (so standard distillation setup) (b) imitability loss with actual student network. If (a), then why is the method “w/o align (ours)” in the result tables? Also, can the authors explain why the teacher returns w/o alignment are lower than with alignment in Figure 4, intuition suggests that it should be higher?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The motivation of the problem and proposed approach is presented intuitively (but lacks some clarity see clarification questions Q1-2).", "weaknesses": "* The paper’s positioning in the related work as the only one considering information asymmetry in student-teacher framework is incorrect. There are several recent works that have looked at this problem [1-5], in fact the imitability cost like the one proposed is explored in [3].\n\n* As a consequence of the above, I think the paper misses key baselines that actually tackle problems under similar assumptions of asymmetry (the considered BC and DAgger baselines are bound to fail here) and therefore does not establish novelty.\n\n* Additionally, the introduced approach involves several hyperparameters – the balancing term for “imitability” cost, the update style (frequency/batch sizes) of the alignment phase. The paper does not describe the implications of these choices of the algorithm and in-reality these can be considerably hard to tune for each task and can raise concerns of applicability and reproducibility of the approach.\n\n_References_:\n\n[1] Warrington, Andrew, et al. \"Robust asymmetric learning in POMDPs.\" International Conference on Machine Learning. PMLR, 2021.\n\n[2] Weihs, Luca, et al. \"Bridging the imitation gap by adaptive insubordination.\" Advances in Neural Information Processing Systems 34 (2021): 19134-19146.\n\n[3] Nguyen, Hai, et al. \"Leveraging fully observable policies for learning under partial observability.\" arXiv preprint arXiv:2211.01991 (2022).\n\n[4] Walsman, Aaron, et al. \"Impossibly good experts and how to follow them.\" The Eleventh International Conference on Learning Representations. 2022.\n\n[5] Shenfeld, Idan, et al. \"TGRL: An algorithm for teacher guided reinforcement learning.\" International Conference on Machine Learning. PMLR, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AgT11PAZ4N", "reviewer": "Reviewer_Xbsv", "summary": "To address the asymmetry in observation spaces during teacher-student training in reinforcement learning, this paper introduces a KL divergence penalty term between the teacher policy and a proxy for the student policy. This penalty is added to the teacher's RL loss, encouraging the teacher to adapt its behavior based on the student's observation space. To facilitate efficient penalty computation during teacher training, the authors employ a proxy student with the same privileged observation space as the teacher. The student learns by imitating the teacher from its limited observation space, while the proxy student imitates the student using privileged observations. They validate their approach in simulation across a toy grid navigation task, a vision-based quadrotor obstacle avoidance environment, and a vision-based manipulation environment, demonstrating the emergence of perception-aware behaviors.", "questions": "- Why does the \"w/o align\" baseline perform so well compared to BC and DAgger, particularly in the quadrotor environment? It seems this baseline would produce behaviors similar to that of BC and DAgger.\n- Does the proposed method require more interactions than the baselines? Were baselines evaluated with equivalent budgets of privileged and high-dimensional samples?\n- What is the performance without using a shared action decoder?\n- How sensitive is the method to the weight of the KL divergence penalty term?\n- Is it possible to have videos of the experiments in simulation ?\n- Out of curiosity, have you experimented with using forward vs reverse KL divergence as the penalty term?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- Impressive qualitative results: The emergence of perception-aware behaviors is impressive, and the authors effectively demonstrate that their approach produces this outcome across different environments, including a quadrotor flight task and a vision-based manipulation task.\n- Interesting approach: The teacher-student observation asymmetry is a relevant problem to address in robotics. The idea of penalizing the teacher for behaviors that cannot be reproduced from the student’s observation space is both simple and sound.", "weaknesses": "- Missing important comparison and discussion: Several prior works focus on training the student with RL to achieve behaviors distinct from the teacher [1, 2, 3, 4, 5, 6, 7] (and likely more). While some of these are referenced in the paper, there is no thorough discussion of this class of methods, and the authors only compare their approach to BC and DAgger, which were not specifically designed to handle observation asymmetry. Since many of these prior works are motivated by teacher-student observation asymmetry, it is crucial that the authors include a proper discussion of this line of work and compare their approach to some of them.\n- Clarity and methodological details: Certain aspects of the approach, especially in Section 4, lack clarity. The method involves alternating between data collection, RL policy updates with the proposed penalty, and policy alignment, but critical details are reserved for the appendix, such as the ratio of privileged to high-dimensional observation samples or the number of updates performed in each phase for each component. A potential drawback of this approach is that alternating training for both teacher and student before teacher training is complete might demand more high-dimensional samples than BC or DAgger, which train the teacher once and then the student. It is unclear how comparisons to baselines were made, particularly regarding environment interactions. Does the proposed method require more interactions than the baselines? Were baselines evaluated with equivalent budgets of privileged and high-dimensional samples? Additionally, compared to the baselines, the approach introduces a proxy student, a more involved training protocol and a KL penalty with associated hyperparameters, making it challenging to assess the algorithm's overall complexity. The “w/o Align” baseline is also not clearly explained, nor is it clear why its behavior differs from BC and DAgger. Finally, the shared action decoder is introduced without sufficient motivation; the authors could clarify this design choice, for instance through ablation studies.\n\nReferences:\n- [1] TGRL: An Algorithm for Teacher Guided Reinforcement Learning, ICML2023\n- [2] Leveraging Fully Observable Policies for Learning under Partial Observability, CoRL2022\n- [3] SoloParkour: Constrained Reinforcement Learning for Visual Locomotion from Privileged Experience, CoRL2024\n- [4] Bootstrapping reinforcement learning with imitation for vision-based agile flight, CoRL2024\n- [5] Privileged Sensing Scaffolds Reinforcement Learning. ICLR2024\n- [6] Real-World Humanoid Locomotion with Reinforcement Learning, 2023\n- [7] Bridging the Imitation Gap by Adaptive Insubordination, NeurIPS2021", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)