Pretraining in Actor-Critic Reinforcement Learning for Robot Locomotion

TL;DR

The paper addresses slow sample efficiency in RL for robot locomotion by introducing an embodiment-aware pretraining paradigm. It trains a Proprioceptive Inverse Dynamics Model (PIDM) on task-agnostic exploration data to warm-start PPO actor-critic networks, enabling faster adaptation to downstream tasks on the same embodiment. Across 9 environments and 3 robot embodiments, the method yields consistent improvements in final performance (~7.3%) and substantial gains in sample efficiency (~36.9%), with thorough ablations clarifying the mechanisms. The approach offers a practical, plug-in initialization that generalizes across tasks without requiring task-specific rewards during pretraining.

Abstract

The pretraining-finetuning paradigm has facilitated numerous transformative advancements in artificial intelligence research in recent years. However, in the domain of reinforcement learning (RL) for robot locomotion, individual skills are often learned from scratch despite the high likelihood that some generalizable knowledge is shared across all task-specific policies belonging to the same robot embodiment. This work aims to define a paradigm for pretraining neural network models that encapsulate such knowledge and can subsequently serve as a basis for warm-starting the RL process in classic actor-critic algorithms, such as Proximal Policy Optimization (PPO). We begin with a task-agnostic exploration-based data collection algorithm to gather diverse, dynamic transition data, which is then used to train a Proprioceptive Inverse Dynamics Model (PIDM) through supervised learning. The pretrained weights are then loaded into both the actor and critic networks to warm-start the policy optimization of actual tasks. We systematically validated our proposed method with 9 distinct robot locomotion RL environments comprising 3 different robot embodiments, showing significant benefits of this initialization strategy. Our proposed approach on average improves sample efficiency by 36.9% and task performance by 7.3% compared to random initialization. We further present key ablation studies and empirical analyses that shed light on the mechanisms behind the effectiveness of this method.

Figures (18)

• Figure 1: Method overview: We (i) collect task-agnostic data using an exploration-driven policy, (ii) to train a Proprioceptive Inverse Dynamics Model (PIDM) to capture embodiment-aware dynamics, and (iii) initialize the actor-critic networks in PPO to warm-start the RL process.
• Figure 2: Diagram for exploration-based data collection pipeline, showcasing how the simulation collects data, and is guided by the ensemble of PIDM models that are periodically retrained using the buffered data.
• Figure 3: Proprioceptive Inverse Dynamics Model (PIDM) architecture and its integration into the actor network. During pretraining of the PIDM, the dashed red parts of the network are included. However, when integrating into the actor-critic structure, those are removed and replaced by the encoder and decoder in purple.
• Figure 4: Experiment to probe dynamics knowledge in vanilla policy networks. We analyze the correlation between intermediate representations between layers and the future joint state $q_{t+1}$. The zero-order extrapolation in (b) is a reference of the accuracy of always predicting $q_t = q_{t+1}$. Shaded areas indicate the standard deviation over 5 RL runs.
• Figure 5: PIDM training and dataset analysis of ANYmal D: For the pretraining dataset we visualize its (a) coverage (green) compared to the initial exploration stages in RL (red and blue) using an UMAP projection and (b) the sample distribution of absolute action magnitudes $|a_t - q_t|$ over different terrains. Finally, in (c) we show the resulting PIDM accuracy across the entire action range as absolute joint errors $|\hat{q}_{t+1} - q_{t+1}|$ and also normalized by the action magnitude.