Noise-conditioned Energy-based Annealed Rewards (NEAR): A Generative Framework for Imitation Learning from Observation

TL;DR

The paper tackles imitation learning from observation by replacing adversarial reward learning with an energy-based framework built on score-based generative models. By applying Noise-conditioned Score Networks to learn a family of smooth energy functions from expert state transitions, NEAR provides stable, well-defined rewards for RL without adversarial min-max training. An annealing strategy gradually shifts the reward landscape to guide policies toward the expert distribution, achieving competitive imitation on complex humanoid tasks and showing advantages in stability and trajectory smoothness. This approach reduces sensitivity to hyperparameters typical of GAN-based IL and demonstrates practical potential for learning from limited, state-only data in physically grounded robots.

Abstract

This paper introduces a new imitation learning framework based on energy-based generative models capable of learning complex, physics-dependent, robot motion policies through state-only expert motion trajectories. Our algorithm, called Noise-conditioned Energy-based Annealed Rewards (NEAR), constructs several perturbed versions of the expert's motion data distribution and learns smooth, and well-defined representations of the data distribution's energy function using denoising score matching. We propose to use these learnt energy functions as reward functions to learn imitation policies via reinforcement learning. We also present a strategy to gradually switch between the learnt energy functions, ensuring that the learnt rewards are always well-defined in the manifold of policy-generated samples. We evaluate our algorithm on complex humanoid tasks such as locomotion and martial arts and compare it with state-only adversarial imitation learning algorithms like Adversarial Motion Priors (AMP). Our framework sidesteps the optimisation challenges of adversarial imitation learning techniques and produces results comparable to AMP in several quantitative metrics across multiple imitation settings.

Figures (9)

• Figure 1: A comparison of reward functions (probability density approximations) learnt in a 2D target-reaching imitation task (left). In this task, an agent aims to reach a goal and expert demonstrations ($p_D$) pass through an L-shaped maze. The learnt reward function is expected to encourage the agent to pass through the maze. In the middle, we show $\texttt{rew}(s'|s)$ for all reachable states around a state $s$ (green circle) at different training epochs. On the right, we show an illustration of the non-smooth reward landscape of adversarial IL. The energy-based reward is a smooth (with continuous gradients), accurate representation of $p_D$ and is constant regardless of the distribution of policy-generated motions ($p_G$). In contrast, the adversarial reward is non-smooth and prone to instability. Additionally, it changes depending on $p_G$ (the discriminator tends to minimise policy predictions) and can provide non-stationary reward signals (additional details in \ref{['appendix:maze_domain']}).
• Figure 2: Degradation of an adversarially learnt policy (AMP) in a stylised walking imitation task. With sufficient training, the policy does learn to complete the task, however, performance fluctuates substantially throughout training (with degradation seen at $16e6$ training samples).
• Figure 3: Annealing (during RL) ensures that the agent always receives a focused and well-defined reward signal, thereby motivating the policy to produce motions similar to $p_D$. Here, $p_D$ is a distribution of the expert's state transitions in a 2D target-reaching task (introduced in \ref{['fig:illustration']}). The learnt energy functions $e_{\theta}(\cdot, \sigma_k)$ are illustrated as dilated versions (L-shaped boundaries) of $p_D$ that are well-defined only inside their respective perturbed manifold ("inner" regions of the L-shaped boundaries). The manifold of policy-generated motions is indicated by $\mathop{\mathrm{supp}}\nolimits(\pi_{\theta_G})$. The policy is shown to have improved from left to right since $\mathop{\mathrm{supp}}\nolimits(\pi_{\theta_G})$ for the improved policy is closer to $p_D$. The reward function currently maximised by the agent is highlighted in red. During RL, the agent starts at the energy function (reward) of a lower noise level ($e_{\theta}(\cdot, \sigma_k)$) and switches to a higher one ($e_{\theta}(\cdot, \sigma_{k+1})$) upon receiving a sufficiently high average return. Arrows indicate the gradient of the rewards in $\mathop{\mathrm{supp}}\nolimits(\pi_{\theta_G})$ (avg. score).
• Figure 4: Snapshots of the policies trained with NEAR. Mummy-style walking and spin-kick are single-clip imitation tasks. The bottom row shows goal-conditioned RL policies that also optimise an environment-provided task reward.
• Figure 5: Annealing at higher noise levels causes a drop in the energy reward's return.

Theorems & Definitions (4)

• Lemma A.1
• proof
• Theorem A.2
• proof