[Paper Review] Domain-Adversarial and Conditional State Space Model for Imitation Learning

Ryo Okumura, Masashi Okada+1 more|arXiv (Cornell University)|Jan 31, 2020

Fault Detection and Control Systems43 references4 citations

TL;DR

This paper proposes a domain-adversarial and -conditional state space model (DAC-SSM) that learns domain-agnostic, task- and dynamics-aware state representations for imitation learning in partially observable environments with large domain shifts. By jointly optimizing state inference, observation reconstruction, forward dynamics, and reward modeling through adversarial training with domain discriminators and domain-conditioned encoders/decoders, DAC-SSM enables agents to achieve expert-level performance in continuous control tasks—outperforming baselines by over 2x in sparse-reward simulators.

ABSTRACT

State representation learning (SRL) in partially observable Markov decision processes has been studied to learn abstract features of data useful for robot control tasks. For SRL, acquiring domain-agnostic states is essential for achieving efficient imitation learning. Without these states, imitation learning is hampered by domain-dependent information useless for control. However, existing methods fail to remove such disturbances from the states when the data from experts and agents show large domain shifts. To overcome this issue, we propose a domain-adversarial and conditional state space model (DAC-SSM) that enables control systems to obtain domain-agnostic and task- and dynamics-aware states. DAC-SSM jointly optimizes the state inference, observation reconstruction, forward dynamics, and reward models. To remove domain-dependent information from the states, the model is trained with domain discriminators in an adversarial manner, and the reconstruction is conditioned on domain labels. We experimentally evaluated the model predictive control performance via imitation learning for continuous control of sparse reward tasks in simulators and compared it with the performance of the existing SRL method. The agents from DAC-SSM achieved performance comparable to experts and more than twice the baselines. We conclude domain-agnostic states are essential for imitation learning that has large domain shifts and can be obtained using DAC-SSM.

Motivation & Objective

• To address the challenge of domain shifts in imitation learning, where expert and agent data differ in appearance or context unrelated to control.
• To learn state representations that are invariant to domain-specific distractions (e.g., robot appearance, camera angles) while preserving task-relevant dynamics.
• To improve imitation learning performance in partially observable Markov decision processes (POMDPs) by disentangling domain-invariant control features from domain-specific noise.
• To develop a unified framework that jointly optimizes state inference, reconstruction, dynamics, and reward modeling for robust imitation.
• To validate that domain-agnostic states are essential for effective imitation learning under large domain shifts, especially in sparse-reward environments.

Proposed method

• DAC-SSM extends a recurrent state space model (RSSM) with domain-adversarial training to minimize domain-specific information in latent states.
• A domain discriminator is trained to classify the domain of latent states, and its loss is minimized via gradient reversal to encourage domain invariance.
• The encoder and decoder are conditioned on domain labels to enable disentangled reconstruction of domain-specific and control-relevant features.
• The model jointly optimizes four objectives: state inference, observation reconstruction, forward dynamics prediction, and reward modeling.
• The reward model is implemented as an optimality discriminator that evaluates trajectory quality based on latent states, with imitation rewards computed from domain-agnostic representations.
• The training objective includes a domain confusion loss (adversarial) and a reconstruction loss conditioned on domain labels, enabling disentanglement of control-relevant from domain-specific features.

Experimental results

Research questions

• RQ1 Can domain-agnostic state representations be learned effectively in the presence of large domain shifts in expert and agent data?
• RQ2 Does adversarial training with domain discriminators improve imitation learning performance when control-irrelevant visual variations exist?
• RQ3 How does conditioning the encoder and decoder on domain labels affect the disentanglement of domain-invariant and task-relevant features?
• RQ4 Can DAC-SSM outperform existing SRL methods in imitation learning for sparse-reward continuous control tasks with domain shifts?
• RQ5 What is the impact of hyperparameter tuning (e.g., domain confusion loss coefficient) on the quality of learned state representations?

Key findings

• DAC-SSM achieved performance comparable to expert policies in all evaluated sparse-reward tasks, outperforming the baseline (PlaNet+D_O) by more than 2x in cumulative return.
• The version of DAC-SSM using dual rewards (DAC/dual) achieved the best performance across tasks, except for Cup-Catch, where imitation rewards were suboptimal due to a short planning horizon.
• When domain adversarial training was removed (DA/dual), performance dropped significantly, indicating that adversarial training is essential for domain invariance.
• The ablation study showed that domain-conditional encoders/decoders alone (DC/dual) performed nearly as well as DAC/dual, but only when combined with adversarial training did the model achieve optimal results.
• Reconstruction results confirmed that DAC-SSM captures control-relevant features (e.g., joint angles) while suppressing domain-specific features (e.g., floor color, object texture), demonstrating successful disentanglement.
• The domain confusion loss coefficient λ had task-dependent optimal values, with λ=3 yielding the best performance on Connector-Insertion, indicating that hyperparameter tuning is crucial for robustness.