Generative Adversarial Imitation Learning (GAIL)¶

GAIL performs imitation learning via an actor–discriminator adversarial game: the policy learns to generate trajectories indistinguishable from expert demonstrations without an explicit reward function.

Components¶

Actor-Critic: ActorCritic outputs actions and estimates \(V(s)\)
Discriminator: Discriminator distinguishes expert from agent pairs \((s,a)\)
Policy optimizer: PPO-style clipped surrogate updates

Theory¶

Min–max objective:

\[ \min_{\pi} \max_{D} \; \mathbb{E}_{(s,a)\sim \pi_E}[\log D(s,a)] + \mathbb{E}_{(s,a)\sim \pi}[\log (1 - D(s,a))] \]

Discriminator-derived pseudo reward (for the actor):

\[ r_D(s,a) = -\log D(s,a) \]

PPO actor update (as implemented):

\[ \mathcal{L}_\text{actor} = -\,\mathbb{E}\Big[ \min\big(r_t A_t,\ \mathrm{clip}(r_t,1-\varepsilon,1+\varepsilon) A_t\big) \Big],\quad r_t = \exp(\log \pi_\theta - \log \pi_{\theta_{\text{old}}}) \]

Advantage via GAE:

\[ \delta_t = r_D(s_t,a_t) + \gamma V(s_{t+1}) - V(s_t),\quad \hat{A}_t = \sum_{l\ge 0} (\gamma\lambda)^l\, \delta_{t+l} \]

Expert data¶

Expect expert_data as an array of shape [N, obs_dim + act_dim]: state concatenated with action.

Training loop¶

Generate agent rollouts \(s_t, a_t \sim \pi\); store log_prob, V(s)
Compute pseudo rewards r_D = -log D([s,a]), then GAE returns/advantages
Update actor/critic with PPO mini-batches
Train the discriminator with BCE: D(fake)=1, D(real)=0
Periodically evaluate the policy and early-stop based on max_reward

Example (LinearLongitudinalF16‑v0)¶

import gymnasium as gym
import numpy as np
from tensoraerospace.agent.gail.model import GAIL
from tensoraerospace.utils import generate_time_period
from tensoraerospace.signals.standard import unit_step

dt = 0.01
tp = generate_time_period(tn=20, dt=dt)
number_time_steps = len(tp)
reference_signals = unit_step(degree=5, tp=tp, time_step=1000, output_rad=True).reshape(1, -1)

env = gym.make('LinearLongitudinalF16-v0',
               number_time_steps=number_time_steps,
               initial_state=[[0],[0],[0]],
               reference_signal=reference_signals,
               use_reward=False,
               state_space=["theta","alpha","q"],
               output_space=["theta","alpha","q"],
               control_space=["ele"],
               tracking_states=["alpha"],)

expert_data = np.load('expert_f16.npy')
agent = GAIL(env, learning_rate=3e-3, max_steps=20, mini_batch_size=16, epochs=4, data=expert_data)
agent.learn(max_frames=5000, max_reward=-1)

# Unified API (wraps learn)
agent.train(num_episodes=250, max_steps=20, max_reward=-1)

Unified training interface¶

GAIL exposes the shared unified train() API from BaseRLModel. Internally it delegates to the legacy learn() method, translating num_episodes * max_steps into a max_frames budget. GAIL-specific options accepted via **kwargs:

max_frames (int): override the computed step budget.
max_reward (float): early-stop threshold for the mean test reward.

Tip

High-quality expert_data is crucial—include demonstrations with varied initial states and maneuvers.

Gymnasium 5-tuple API

This implementation uses the modern Gymnasium 5-tuple step API internally:

next_state, reward, terminated, truncated, info = env.step(action)
done = terminated or truncated

If you are migrating from older code that used the 4-tuple API (next_state, reward, done, info = env.step(action)), ensure your environment is compatible with Gymnasium and returns the 5-tuple.

API reference¶

`GAIL(env, learning_rate, max_steps, mini_batch_size, epochs, data, log_dir=None, wandb_project=None, wandb_entity=None, wandb_run_name=None, wandb_tags=None, wandb_config=None, device=None)` ¶

Generative Adversarial Imitation Learning trainer.

Initialize the GAIL algorithm.

Parameters:

Name	Type	Description	Default
`env`	`Env`	Environment instance.	required
`learning_rate`	`float`	Learning rate.	required
`max_steps`	`int`	Maximum steps per rollout/episode.	required
`mini_batch_size`	`int`	Mini-batch size.	required
`epochs`	`int`	Number of training epochs.	required
`data`	`Array`	Expert demonstrations (states/actions).	required
`log_dir`	`Union[str, Path, None]`	Optional TensorBoard log directory. When provided, a `MetricWriter` is created and canonical metrics are emitted during :meth:`learn`. When `None` (default), no logging is performed.	`None`
`device`	`Optional[Union[str, device]]`	Torch device for the networks. Defaults to `"cuda"` when available, otherwise `"cpu"`.	`None`

`expert_reward(state, action)` ¶

Compute imitation reward using the discriminator.

`test_env()` ¶

Run one evaluation rollout and return total reward.

`ppo_update(ppo_epochs, mini_batch_size, states, actions, log_probs, returns, advantages, clip_param=0.2, env_step=0)` ¶

PPO update function.

Parameters:

Name	Type	Description	Default
`ppo_epochs`	`int`	Number of epochs.	required
`mini_batch_size`	`int`	Mini-batch size.	required
`states`	`Tensor`	Batch of states.	required
`actions`	`Tensor`	Batch of actions.	required
`log_probs`	`Tensor`	Batch of action log probabilities.	required
`returns`	`Tensor`	Batch of discounted rewards.	required
`advantages`	`Tensor`	Batch of advantage function values.	required
`clip_param`	`float`	Clipping constant.	`0.2`
`env_step`	`int`	Cumulative environment-step counter used as the x-axis for any TensorBoard scalars written from inside this update. Defaults to 0.	`0`

`train(num_episodes=100, *, max_steps=None, save_best=False, save_path=None, verbose=True, **kwargs)` ¶

Train GAIL (unified interface wrapper around :meth:learn).

Parameters:

Name	Type	Description	Default
`num_episodes`	`int`	Number of episodes. Combined with `max_steps` to produce a `max_frames` budget unless `max_frames` is supplied explicitly via `**kwargs`.	`100`
`max_steps`	`Optional[int]`	Max steps per episode. Defaults to `self.max_steps` when not provided.	`None`
`save_best`	`bool`	Reserved for API consistency.	`False`
`save_path`	`Optional[str]`	Reserved for API consistency.	`None`
`verbose`	`bool`	Reserved for symmetry.	`True`
`**kwargs`	`Any`	GAIL-specific options: `max_frames` (int): override the computed budget. `max_reward` (float, default `inf`): early-stop threshold on mean test reward.	`{}`

Returns:

Name	Type	Description
`dict`	`dict`	Empty dict (GAIL does not yet collect metrics).

`learn(max_frames, max_reward)` ¶

Agent training function.

Parameters:

Name	Type	Description	Default
`max_frames`	`int`	Maximum number of steps in the environment.	required
`max_reward`	`int`	Reward threshold for stopping training.	required

`get_param_env()` ¶

Return serializable configuration of environment and policy.

`save(path=None, save_gradients=False)` ¶

Save GAIL model to the specified directory.

Saves actor-critic network, discriminator network, and configuration. Optionally saves optimizer states for resuming training.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path \| None`	Base save directory. If None, saves to the current working directory.	`None`
`save_gradients`	`bool`	If True, also save optimizer state dicts.	`False`

Returns:

Name	Type	Description
`Path`	`Path`	The directory where the model was saved.

`from_pretrained(repo_name, env=None, data=None, access_token=None, version=None, load_gradients=False)` `classmethod` ¶

Load pretrained model from a local directory or Hugging Face Hub.

Parameters:

Name	Type	Description	Default
`repo_name`	`str`	Path to a local folder or a Hugging Face repo id (e.g. `"namespace/repo_name"`).	required
`env`	`Optional[Env]`	Environment instance (required).	`None`
`data`	`Optional[ndarray]`	Expert demonstration data (optional).	`None`
`access_token`	`Optional[str]`	Hugging Face access token for private repos.	`None`
`version`	`Optional[str]`	Revision / branch / tag on Hugging Face.	`None`
`load_gradients`	`bool`	Restore optimizer states for continued training.	`False`

Returns:

Name	Type	Description
`GAIL`	`GAIL`	Initialized agent.

`publish_to_hub(repo_name, folder_path, access_token=None)` ¶

Upload a saved model folder to Hugging Face Hub.

Parameters:

Name	Type	Description	Default
`repo_name`	`str`	Repository id on Hugging Face (e.g. `"user/my-gail"`).	required
`folder_path`	`Union[str, Path]`	Local folder produced by :meth:`save`.	required
`access_token`	`Optional[str]`	Hugging Face access token.	`None`

References¶

Generative Adversarial Imitation Learning

Tested on¶

Unity environment
LinearLongitudinalF16‑v0 (repository example)

Generative Adversarial Imitation Learning (GAIL)¶

Components¶

Theory¶

Expert data¶

Training loop¶

Example (LinearLongitudinalF16‑v0)¶

Unified training interface¶

API reference¶

GAIL(env, learning_rate, max_steps, mini_batch_size, epochs, data, log_dir=None, wandb_project=None, wandb_entity=None, wandb_run_name=None, wandb_tags=None, wandb_config=None, device=None) ¶

expert_reward(state, action) ¶

test_env() ¶

ppo_update(ppo_epochs, mini_batch_size, states, actions, log_probs, returns, advantages, clip_param=0.2, env_step=0) ¶

train(num_episodes=100, *, max_steps=None, save_best=False, save_path=None, verbose=True, **kwargs) ¶

learn(max_frames, max_reward) ¶

get_param_env() ¶

save(path=None, save_gradients=False) ¶

from_pretrained(repo_name, env=None, data=None, access_token=None, version=None, load_gradients=False) classmethod ¶

publish_to_hub(repo_name, folder_path, access_token=None) ¶

References¶

Tested on¶

`GAIL(env, learning_rate, max_steps, mini_batch_size, epochs, data, log_dir=None, wandb_project=None, wandb_entity=None, wandb_run_name=None, wandb_tags=None, wandb_config=None, device=None)` ¶

`expert_reward(state, action)` ¶

`test_env()` ¶

`ppo_update(ppo_epochs, mini_batch_size, states, actions, log_probs, returns, advantages, clip_param=0.2, env_step=0)` ¶

`train(num_episodes=100, *, max_steps=None, save_best=False, save_path=None, verbose=True, **kwargs)` ¶

`learn(max_frames, max_reward)` ¶

`get_param_env()` ¶

`save(path=None, save_gradients=False)` ¶

`from_pretrained(repo_name, env=None, data=None, access_token=None, version=None, load_gradients=False)` `classmethod` ¶

`publish_to_hub(repo_name, folder_path, access_token=None)` ¶