Global State Encoder in MAIL

• The Global State Encoder (GSE) is a neural module that compresses complex, high-dimensional state observations into fixed latent vectors used by both the policy and discriminator.
• It employs a CNN-based architecture with behavior cloning pre-training, effectively balancing modalities and reducing non-stationarity in adversarial imitation settings.
• Integration into the MAIL framework yields stable training and enhanced performance on tasks like grid navigation and LunarLander-v2 through robust representation sharing.

The Global State Encoder (GSE) is a high-dimensional, shared neural mapping central to the Mature GAIL (MAIL) algorithm for imitation learning, as introduced in "Mature GAIL: Imitation Learning for Low-level and High-dimensional Input using Global Encoder and Cost Transformation" (Shin et al., 2019). The GSE is explicitly designed to compress complex state observations, such as stacked image frames, into a fixed-dimensional latent space utilized by both the policy (generator) and the discriminator (cost network). This symmetrical usage addresses core issues of modality imbalance and non-stationarity inherent in adversarial imitation settings, particularly for environments with high-dimensional state spaces.

1. Definition and Motivation

The Global State Encoder (GSE) in MAIL is defined as a parameterized neural module $E_{\theta_E}: \mathcal{S} \rightarrow \mathbb{R}^d$ mapping from state space $\mathcal{S}$ (potentially an image stack) into a $d$-dimensional latent vector $z$. This encoder is explicitly shared between both the policy network $\pi$ and the discriminator $D$, so that both networks base decision-making and credit assignment on an identical compact representation of state. This design is in contrast with previous methods in Generative Adversarial Imitation Learning (GAIL), which either input raw states or implemented an independent state encoder solely within the discriminator.

The GSE serves three primary functions:

• Dimensionality Reduction: Compresses raw, high-dimensional inputs (e.g., $32\times32\times4$ image stacks) into a lower-dimensional, information-rich latent vector.
• Modality Balancing: Eliminates the imbalance noted by Atrey et al. (2010), where discriminator gradients can dominate due to independent preprocessing pipelines.
• Stability in Adversarial Training: Ensures that both policy and discriminator co-evolve with a consistent, non-drifting state representation, thereby reducing non-stationarity during adversarial min–max optimization.

2. Layerwise Architecture and Instantiation

Architecturally, the GSE is described as a convolutional neural network (CNN) model transferred from a Behavior Cloning (BC) pre-training stage. The precise channel and filter parameters are not given, but alignment with standard low-level visual encoders for similar tasks is implied. A typical instantiation compatible with the data could be summarized as follows:

Layer	Output Channels / Units	Kernel / Stride
Conv₁	32	$8\times8$ / stride 4
Conv₂	64	$4\times4$ / stride 2
Conv₃	64	$3\times3$ / stride 1
Flatten + FC	$\mathcal{S}$0 ($\mathcal{S}$1)	--

The input is an image stack $\mathcal{S}$2; in the navigation benchmark, $\mathcal{S}$3 and $\mathcal{S}$4. The output of the fully-connected layer is interpreted as the latent state $\mathcal{S}$5. No recurrence or explicit attention mechanisms are present. Any normalization layers are absorbed from the BC encoder, and the GSE is fixed (frozen weights) during the main adversarial imitation phase.

3. Mathematical Encoding Process

Given a time-indexed stack of the last $\mathcal{S}$6 frames $\mathcal{S}$7, the GSE computes the latent state through a deterministic mapping:

$\mathcal{S}$8

Equivalently, in the notation of the source, $\mathcal{S}$9, where $d$0 denotes the encoder portion of the pre-trained actor network. This mapping is fully convolutional except for the terminal fully-connected layer and retains no explicit temporal aggregation beyond the spatial stacking of frames as channels.

4. Integration into the Adversarial Imitation Learning Objective

Within the MAIL framework, the encoded latent $d$1 replaces the state $d$2 in both the discriminator and the policy objective. The adversarial loss, adapted from classic GAIL, becomes:

$d$3

A central innovation is the reward penalization mechanism. The reward signal used for policy optimization is given by:

$d$4

This transformation ensures that the reward is zero when the discriminator is maximally confused ($d$5), positive for expert-like behavior, and negative otherwise, addressing reward sparsity and stabilization failures from prior work.

5. Training Regime and Gradient Dynamics

Training with the GSE in MAIL is divided into two sequential phases:

• Phase 1: Behavior Cloning Pre-training
  • Train a standard actor-critic network via supervised behavior cloning on expert demonstrations. The actor is factored as $d$6, with $d$7.
  • Both encoder $d$8 and classifier head $d$9 are updated. Upon convergence, $z$0 is frozen; only $z$1 is discarded.
• Phase 2: Adversarial Imitation (MAIL)
  • Initialize policy $z$2 and discriminator $z$3 with $z$4 weights transferred and fixed.
  • Alternate between: (1) updating $z$5 to maximize the adversarial loss, and (2) updating $z$6 via PPO with the penalized reward. During these updates, gradients do not flow through $z$7, which remains fixed. No additional regularization beyond standard PPO clipping is introduced.

This approach suppresses encoder-induced non-stationarity, as only the policy and discriminator adapt; the global latent representation remains identical throughout adversarial updates.

6. Empirical Effects and Ablation Findings

MAIL is empirically benchmarked on a 7×7 grid hierarchical navigation task (pixel inputs) and OpenAI Gym's LunarLander-v2 (low-dimensional continuous states). Key results include:

• Hierarchical Navigation:
  • Plain GAIL fails to solve the task (best $z$8).
  • GAIL with only GSE (GAIL_GE) or only reward penalization (GAIL_LS) both fail or exhibit instability.
  • MAIL (GSE + reward penalization) reliably solves the task, reaching mean final score $z$9 within 13,000 episodes.
  • MAIL + VDB (adding a Variational Discriminator Bottleneck) further reduces required episodes to 3,000 with a near-optimal score.
• LunarLander-v2:
  • Penalization schemes (scaled, log-shifted, linear, tanh) consistently raise mean returns well above baseline, with tanh penalization yielding $\pi$0.

Ablations demonstrate that only the conjunction of a pre-trained and fixed GSE and the shifted-log reward recovers the desired stability and sample efficiency; all other encoder training regimens lead to collapse or failed learning.

7. Implications and Role in Imitation Learning Research

The introduction of the Global State Encoder marks a shift in adversarial imitation learning protocols toward joint, fixed representation sharing between policy and discriminator. This stabilizes learning for high-dimensional, visual-based tasks previously refractory to GAIL-style solutions. The explicit two-phase (BC pre-training + fixed encoder) methodology addresses both representational drift and gradient pathologies. The approach generalizes readily to other GAIL variants and is further compatible with regularization techniques such as the Variational Discriminator Bottleneck.

A plausible implication is the broader utility of global, fixed encoders across other adversarial RL and IL paradigms where representational coupling and non-stationarity traditionally impede sample efficiency and convergence (Shin et al., 2019).