Problem Space Transformations for Out-of-Distribution Generalisation in Behavioural Cloning

TL;DR

This paper explores how problem space transformations based on robotic manipulation properties like pose equivariance and locality can improve out-of-distribution generalisation of behavioural cloning policies in robotic tasks, demonstrated through simulations and real-world experiments.

Contribution

It introduces the use of problem space transformations derived from manipulation properties to enhance OOD generalisation in behavioural cloning, applicable to both MLP and diffusion-based policies.

Findings

Transformations improve OOD generalisation in simulated environments.

Transformations enhance real-world robotic manipulation performance.

Empirical evidence supports the effectiveness of problem space transformations.

Abstract

The combination of behavioural cloning and neural networks has driven significant progress in robotic manipulation. As these algorithms may require a large number of demonstrations for each task of interest, they remain fundamentally inefficient in complex scenarios, in which finite datasets can hardly cover the state space. One of the remaining challenges is thus out-of-distribution (OOD) generalisation, i.e. the ability to predict correct actions for states with a low likelihood with respect to the state occupancy induced by the dataset. This issue is aggravated when the system to control is treated as a black-box, ignoring its physical properties. This work highlights widespread properties of robotic manipulation, specifically pose equivariance and locality. We investigate the effect of the choice of problem space on OOD performance of BC policies and how transformations arising from characteristic properties of manipulation can be employed for its improvement. Through controlled, simulated and real-world experiments, we empirically demonstrate that these transformations allow behaviour cloning policies, using either standard MLP-based one-step action prediction or diffusion-based action-sequence prediction, to generalise better to certain OOD problem instances. Code is available at https://github.com/kirandoshi/pst_ood_gen.