PI-ARS: Accelerating Evolution-Learned Visual-Locomotion with Predictive Information Representations

Authors: Kuang-Huei Lee, Ofir Nachum, Tingnan Zhang, Sergio Guadarrama, Jie Tan, Wenhao Yu · IROS 2022

Licences

Open Access
Paper: CC BY 4.0

Everything below draws on the full paper under its CC BY 4.0 licence; its figure is credited to the authors.

TL;DR

A 128-d depth and proprioception encoder trained by predictive information, with a small policy head trained by ARS: 100% real stepping-stone success against 40% for ARS.

Why it matters

It keeps the parallelism and simplicity of evolution strategies while shrinking their main obstacle for vision-based control, the size of the parameter search space. The policies transferred from PyBullet to a real Laikago for stepping stones and indoor navigation using the same sim-to-real procedure as the earlier ARS work.

Problem

Evolution strategies such as Augmented Random Search (ARS) parallelise well and explore in parameter space, but scale poorly to the many parameters of the convolutional networks needed for visual locomotion. In the authors' tests the gradient-based alternatives SAC and PI-SAC made no meaningful progress on these tasks, and the earlier ARS visual-locomotion policy reached only 40% success on real stepping stones.

Proposed solution

PI-ARS splits the agent into an encoder and a small policy head. The encoder maps two 32×24 depth images and proprioceptive states to a 128-d vector and is trained by gradient descent to maximise predictive information, combining an InfoNCE term that predicts the future depth observation with RNN reward prediction, on trajectories that ARS already collects; ARS then optimises only a 3-layer MLP head that outputs high-level commands to an MPC-based low-level controller.

Key insights

Separating representation learning (gradient-based) from policy search (gradient-free) avoids the scaling problem of evolution strategies while keeping their advantages.
Trajectories from ARS perturbation rollouts are stored in a replay buffer and reused to train the encoder, so representation learning needs no extra environment data.
Reward prediction alone gave representations too weak to solve the tasks; the contrastive future-observation term was needed.
The authors suggest ARS's parameter-space exploration and stability are critical here: adding predictive information to SAC (PI-SAC) did not help, so a good representation alone was not enough.

Main results

Real stepping stones (four stones, three gaps of 0.12–0.18 m): PI-ARS succeeds in 10 of 10 trials (100%), against 40% for the ARS baseline under the same training and evaluation setting.
Simulated uneven stepping stones: mean return after 2,000,000 training episodes rises 48.01%, from 2.93 to 4.34.
PI-ARS gives higher returns and better sample efficiency than ARS on all four simulated tasks: uneven stepping stones, quincuncial piles, moving platforms and indoor navigation.
SAC and PI-SAC failed to make any non-negligible learning progress on these tasks despite extensive hyperparameter tuning.
Real indoor navigation: the policy reached the target and found a shortcut between two obstacles, though it collided with one obstacle's arm.
Training used 1024 perturbations per ARS iteration with the top 512 kept, up to 4000 iterations (at most 2,048,000 simulation episodes per trial), and 30 trials per method.

Limitations

In real indoor navigation the robot collided with an obstacle's arm because it was trained only on box-shaped simulated obstacles.
Only indoor tasks were tested; outdoor visual locomotion is left to future work.
The real stepping-stone policy was trained on an easier version of the simulated terrain, with smaller changes in stone height.
Indoor navigation needed a motion capture system to track the robot's base position for the relative goal vector.
The learned policy outputs high-level commands only and depends on an MPC-based low-level controller to execute them.

Field impact

Finalist for the IROS 2022 Best Paper Award and ABB Best Student Paper Award, with 21 citations. The authors publish a supplementary video page; the paper does not state a code release.

Abstract

Evolution Strategy (ES) algorithms have shown promising results in training complex robotic control policies due to their massive parallelism capability, simple implementation, effective parameter-space exploration, and fast training time. However, a key limitation of ES is its scalability to large capacity models, including modern neural network architectures. In this work, we develop Predictive Information Augmented Random Search (PI-ARS) to mitigate this limitation by leveraging recent advancements in representation learning to reduce the parameter search space for ES. Namely, PI-ARS combines a gradient-based representation learning technique, Predictive Information (PI), with a gradient-free ES algorithm, Augmented Random Search (ARS), to train policies that can process complex robot sensory inputs and handle highly nonlinear robot dynamics. We evaluate PI-ARS on a set of challenging visual-locomotion tasks where a quadruped robot needs to walk on uneven stepping stones, quincuncial piles, and moving platforms, as well as to complete an indoor navigation task. Across all tasks, PI-ARS demonstrates significantly better learning efficiency and performance compared to the ARS baseline. We further validate our algorithm by demonstrating that the learned policies can successfully transfer to a real quadruped robot, for example, achieving a 100% success rate on the real-world stepping stone environment, dramatically improving prior results achieving 40% success.

Open Source Robotics