Source-linked AI summary

Object-Aware Regularization for Addressing Causal Confusion in Imitation Learning

Jongjin Park, Younggyo Seo, Chang Liu, Li Zhao, Tao Qin, Jinwoo Shin, Tie-Yan Liu

arXiv:2110.14118v1cs.LGcs.AI

TL;DR

Behavioral cloning can rely on nuisance variables correlated with expert actions rather than their true causes. OREO uses VQ-VAE codes to identify semantic objects and masks same-code feature units during policy training. The paper reports improved BC performance and comparisons favoring OREO over other regularization and causality-based methods.

  • Problem

    Behavioral cloning can exploit nuisance variables correlated with expert actions instead of their true causes, creating a causal-confusion challenge.

    Source: §1 · p. 1
  • Method

    OREO uses VQ-VAE discrete codes to group semantic-object units and randomly drops same-code units while training the imitation policy.

    Source: §1 · p. 2, §3.2 · p. 5
  • Results

    OREO significantly improves BC on confounded Atari environments, achieving mean HNS of 105.6% versus DropBlock’s 91.7%.

    Source: §4 · p. 7
  • Takeaways & Limitations

    OREO improves behavioral cloning without costly environment interaction, a setting the paper identifies as important for safe and successful imitation learning.

    Source: §5 · p. 10

    OREO is designed for image inputs and is not applicable to state-based environments; it regularizes against nuisance correlates rather than deducing cause-effect relations.

    Source: §5 · p. 10

Abstract

from arXiv · show

Behavioral cloning has proven to be effective for learning sequential decision-making policies from expert demonstrations. However, behavioral cloning often suffers from the causal confusion problem where a policy relies on the noticeable effect of expert actions due to the strong correlation but not the cause we desire. This paper presents Object-aware REgularizatiOn (OREO), a simple technique that regularizes an imitation policy in an object-aware manner. Our main idea is to encourage a policy to uniformly attend to all semantic objects, in order to prevent the policy from exploiting nuisance variables strongly correlated with expert actions. To this end, we introduce a two-stage approach: (a) we extract semantic objects from images by utilizing discrete codes from a vector-quantized variational autoencoder, and (b) we randomly drop the units that share the same discrete code together, i.e., masking out semantic objects. Our experiments demonstrate that OREO significantly improves the performance of behavioral cloning, outperforming various other regularization and causality-based methods on a variety of Atari environments and a self-driving CARLA environment. We also show that our method even outperforms inverse reinforcement learning methods trained with a considerable amount of environment interaction.

1 Introduction

Behavioral cloning learns from expert demonstrations but can rely on nuisance variables correlated with expert actions rather than their true causes. OREO addresses this causal confusion by masking semantic objects, and the paper reports improvements over other regularization and causality-based methods.

1 Introduction

  • BC can exploit nuisance variables correlated with expert actions instead of learning their true causes, a problem that can undermine stable deployment.In Pong, policies trained with visible scores struggled when scores were masked, while policies trained with masked scores generalized to the original images.Source: §1 · p. 1
  • Causal discovery is difficult for high-dimensional inputs and may require costly environment interaction, motivating methods that address causal confusion without interaction.The paper notes that observational causal discovery is generally impossible without certain conditions and that image inputs may not be structured as variables in a causal graph.Source: §1 · p. 2
  • OREO extracts semantic objects using VQ-VAE discrete codes, then jointly drops feature-map units sharing a code to regularize the policy.The method aims to encourage attention across semantic objects rather than reliance on nuisance correlates.Source: §1 · p. 2
  • The paper reports that OREO improves BC on confounded Atari environments, outperforming other regularization and causality-based methods.It also reports outperforming inverse reinforcement learning methods trained with considerable environment interaction.Source: §1 · p. 2

2 Related work

Related work frames imitation learning around inverse reinforcement learning and behavioral cloning, and reviews efforts to address causal confusion arising alongside distributional shift. The paper positions OREO as policy regularization for high-dimensional inputs without environment interaction.

2 Related work

  • Inverse reinforcement learning identifies a cost function under which the expert is uniquely optimal, while behavioral cloning predicts expert actions from states.The paper uses behavioral cloning, which avoids costly and dangerous environment interaction.Source: §2 · pp. 2–3
  • Causal discovery methods relying on the faithfulness condition are not applicable to imitation learning setups with nuisance correlates [12].The cited work reports that the condition does not hold in those environments.Source: §2 · p. 2
  • Prior approaches address distributional shift through expert queries or regularization, while causal-confusion methods include dropping prior samples and adversarially removing information about prior actions.de Haan et al. [12] use randomly masked disentangled representations and infer the best mask through environment interaction; OREO instead regularizes without that interaction.Source: §2 · p. 3
  • Causal discovery from observational data often assumes disentangled input variables, an assumption that may not hold for high-dimensional images.Although observational causal discovery from images has been demonstrated, combining causal models with representation learning in such domains remains an open problem [46].Source: §2 · p. 3

3 Method

The method sets up imitation learning from expert demonstrations without environment interaction, then combines behavioral cloning with object-aware masking. OREO uses VQ-VAE codes to group feature-map units and drops same-code units together during policy training.

3.1 Preliminaries

  • The setup treats imitation learning as learning a policy from expert demonstrations in a discrete-time MDP without an explicitly defined reward.In this problem setting, the agent cannot interact with the environment and learns using demonstrations alone.Source: §3.1 · p. 3
  • Behavioral cloning trains a policy and convolutional encoder to predict expert actions from states by minimizing their negative log-likelihood.The policy is modeled as a multinomial distribution over actions for discrete action spaces.Source: §3.1 · p. 4
  • A VQ-VAE encodes images into latent vectors, quantizes them using a learned codebook, and trains a decoder to reconstruct the images from the discrete representations.Its objective includes reconstruction, codebook, and commitment losses.Source: §3.1 · p. 4

3.2 OREO: Object-aware regularization for behavioral cloning

  • OREO first trains a VQ-VAE on demonstration images, then regularizes the policy by randomly dropping feature-map units that share discrete codes.The method is intended to encourage the policy to attend across semantic objects rather than exploit nuisance correlates.Source: §3.2 · p. 4
  • OREO treats units mapped to the same or similar VQ-VAE codes as semantic objects and drops same-code units together during training.The mask is generated by sampling a binary variable for each code, with drop probability p, and applying the corresponding code-based mask to the feature map.Source: §3.2 · p. 5
  • For policy training, OREO initializes the convolutional encoder from the VQ-VAE encoder and minimizes expert-action negative log-likelihood on masked features.The mask is shared across feature-map channels, and masked features are scaled by 1/(1 − p) during training.Source: §3.2 · p. 5

4 Experiments

The experiments test OREO against regularization, causality-based, and inverse reinforcement learning methods on confounded Atari environments, and evaluate it on CARLA. They also examine validation, hyperparameters, demonstration size, and how OREO affects visual attention.

Comparative evaluation

  • OREO achieved mean HNS of 105.6% versus 91.7% for DropBlock, the second-best method, and the best scores on 15 of 27 Atari environments.OREO also achieved the best median HNS across the environments.Source: Table 1 · p. 6, §4 · p. 7

Comparison with inverse reinforcement learning methods

  • OREO outperformed GAIL and DRIL on most confounded Atari environments without interaction, though IRL methods might surpass it asymptotically with more interaction.The IRL methods were trained with 20M environment steps.Source: §4 · p. 7, Figure 5 · p. 7

Why is regularization necessary in confounded environments?

  • On confounded Pong, validation accuracy did not align with test score, as validation data could differ from the policy-induced distribution.The authors note that training-time environment evaluation may be dangerous or impossible, motivating policy regularization.Source: §4 · p. 8, Figure 6 · p. 8

Effects of hyperparameters

  • Performance improved as drop probability increased, while codebook sizes that were too small or large could hurt; reported defaults were p = 0.5 and K = 512.The authors note that further hyperparameter tuning could improve OREO.Source: §4 · p. 8, Figure 6 · p. 8

Effects of expert demonstration size

  • Across demonstration sizes N = 5–50, OREO consistently improved BC and beat Dropout and DropBlock except at N = 5, where limited data could destabilize VQ-VAE learning.DropBlock and OREO also consistently outperformed Dropout.Source: §4 · p. 8, Figure 7 · p. 8

Contribution of a separate convolutional encoder

  • OREO reached mean HNS of 100.1% versus 53.0% for VQ-VAE + DropBlock; fixed VQ-VAE representations lacked fine-grained features needed to imitate actions.The comparison shows OREO is not simply BC on fixed VQ-VAE representations.Source: Table 2 · p. 9, §4 · p. 9

How does OREO improve the performance of BC?

  • OREO's encoder attended to important objects such as approaching cars, while BC focused on a small region around a car and scoreboard.The attention-map comparison was made in confounded and original Enduro environments.Source: Figure 8 · p. 9, §4 · p. 9

Effectiveness on real-world applications

  • On daytime CARLA, OREO improved BC and achieved the best success rate on all four benchmark tasks.The evaluation used 150 expert demonstrations and four-run averages.Source: Table 3 · p. 10, §4 · p. 10

5 Discussion

OREO regularizes image-based imitation policies by dropping feature-map units that share VQ-VAE codes. The discussion notes its scope limits and potential privacy, compute, and safety concerns.

5 Discussion

  • OREO randomly drops feature-map units sharing VQ-VAE discrete codes and improves behavioral cloning without costly environment interaction.The method is designed for image inputs and does not infer cause-effect relations; it regularizes policies against exploiting nuisance correlates.Source: §5 · p. 10
  • OREO is designed for image inputs, not state-based environments, and regularizes against nuisance correlates rather than deducing cause-effect relations.The authors note that learning structured disentangled variables and causal relations remains an open problem.Source: §5 · p. 10
  • VQ-VAE use may expose training data to privacy attacks and add computing and energy costs; policies can also imitate dangerous actions in demonstrations.The authors identify differentially private VQ-VAE as necessary for addressing privacy vulnerability in real applications.Source: §5 · p. 10

B.1 Experimental setup

The experiments use expert Atari demonstrations and evaluate final policies on confounded environments. The implementation details cover OREO, regularization and causality-based baselines, and interaction-based IRL methods.

B.1 Experimental setup

  • Experiments use DQN Replay demonstrations from 27 Atari environments, with N ∈ {20, 50}, 84 × 84 grayscale images, and confounding by adding previous actions.Eight environments were selected for Figure 7; Table 4 reports dataset sizes by environment.Source: §B.1 · p. 15, Table 4 · p. 15
  • Without interaction, policies train for 1000 epochs and are scored over 100 episodes; GAIL and DRIL are evaluated every 1M interaction steps over 10 episodes.Policies are evaluated on confounded environments, and the reported score averages episodes for each random seed.Source: §B.1 · p. 15

B.2 Implementation details

  • OREO precomputes VQ-VAE codes and uses five random masks per image; its VQ-VAE uses a 512-code book and is trained for 1000 epochs.Aggregating losses across multiple masks marginally improves performance; training uses batch size 1024 and learning rate 3e-4.Source: §B.2 · p. 16
  • Regularization baselines include four-pixel-padded random crops, Dropout with p = 0.5, DropBlock with p = 0.3 and size 3, and 10 × 10 to 30 × 30 Cutout patches.DropBlock’s probability is linearly increased during training.Source: §B.2 · p. 16
  • The CCIL reproduction adds action prediction to β-VAE training; the CRLR comparison uses a categorical version operating on VQ-VAE codes.The authors report that action prediction was crucial for raising CCIL policy accuracy above chance.Source: §B.2 · p. 16
  • GAIL and DRIL use environment interaction, respectively learning from a discriminator reward or an ensemble-policy disagreement cost.The reported IRL results use a public implementation modified for confounded Atari environments.Source: §B.2 · pp. 16–17

C Comparative evaluation on original Atari environments

On original Atari environments without environment interaction, OREO improves behavioral cloning and leads the compared methods on aggregate human-normalized scores.

C Comparative evaluation on original Atari environments

  • OREO significantly improves behavioral cloning and outperforms the other compared methods on original Atari environments without interaction.The authors connect this result to causal confusion that naturally occurs in the dataset.Source: C Comparative evaluation on original… · p. 17
  • 114.9% mean HNS versus DropBlock’s 99.0%, with OREO best on 14/27 environments and best median and mean HNS overall.The results for each environment are mean returns across eight runs.Source: Table 5 · p. 17, C Comparative evaluation on original… · p. 17
  • The reported comparison covers policies trained on original Atari environments, rather than confounded environments.Table 5’s caption specifies the evaluation setting and reports mean returns across eight runs.Source: Table 5 · p. 17

D CCIL with environment interaction

With environment interaction, CCIL improves, but OREO still outperforms it in most of the six confounded Atari environments compared.

D CCIL with environment interaction

  • CCIL’s environment interaction uses targeted intervention, yet the authors note the difficulty of learning disentangled representations from high-dimensional images.The section reports that CCIL improves during 100 episodes of environment interaction.Source: D CCIL with environment interaction · p. 18
  • OREO outperforms CCIL in most of six confounded Atari environments, although interaction slightly improves CCIL’s performance.Figure 9 compares eight-run means and standard deviations; CCIL† denotes results without environment interaction.Source: Figure 9 · p. 18, D CCIL with environment interaction · p. 18

E Applying OREO to inverse reinforcement learning

The section applies OREO to DRIL’s behavioral-cloning components and reports improved sample-efficiency on confounded CrazyClimber and Pong.

E Applying OREO to inverse reinforcement learning

F OREO with a sequence of observations

The section extends OREO to four stacked observations by masking matching discrete codes across frames. It reports improved BC performance on confounded Atari environments.

F OREO with a sequence of observations

G Comparison with DropBottleneck

The section compares OREO with DropBottleneck and describes a categorical CRLR variant for image-based features. It reports OREO’s Atari results and discusses when feature removal addresses causal confusion.

G Comparison with DropBottleneck

  • DropBottleneck with expert actions is comparable to OREO in some environments, but OREO significantly outperforms it in most; the state-target variant is no better than BC in most environments.CrazyClimber is an exception for both comparisons; Alien, KungFuMaster, and Pong are examples where OREO outperforms the expert-action variant.Source: G Comparison with DropBottleneck · p. 19
  • Removing task-irrelevant information may not address causal confusion because both confounders and true causes can be task-relevant and strongly correlated with expert actions.Source: G Comparison with DropBottleneck · p. 19

H A categorical version of CRLR

  • The categorical CRLR variant one-hot encodes categorical features and balances confounders for pairs of categorical variables; on images, it uses VQ-VAE discrete codes.The policy and sample weights are jointly trained with a weighted BC objective and the proposed regularizer.Source: H A categorical version of CRLR · p. 20

I Extended experimental results on confounded Atari environments

Without environment interaction, OREO achieves the best scores on 14 of 27 original Atari environments and the highest median and mean human-normalized scores.

I Extended experimental results on confounded Atari environments

  • Each environment’s reported return is the mean and standard deviation over eight runs.Source: Table 10
Loading 2110.14118v1…