<jupyter_start><jupyter_text>Train an Agent using Adversarial Inverse Reinforcement Learning As usual, we first need an expert. Note that we now use a variant of the CartPole environment from the seals package, which has fixed episode durations. Read more about why we do this [here](https://imitation.readthedocs.io/en/latest/guide/variable_horizon.html).<jupyter_code>from stable_baselines3 import PPO
from stable_baselines3.ppo import MlpPolicy
import gym
import seals

env = gym.make("seals/CartPole-v0")
expert = PPO(
    policy=MlpPolicy,
    env=env,
    seed=0,
    batch_size=64,
    ent_coef=0.0,
    learning_rate=0.0003,
    n_epochs=10,
    n_steps=64,
)
expert.learn(1000)  # Note: set to 100000 to train a proficient expert<jupyter_output><empty_output><jupyter_text>We generate some expert trajectories, that the discriminator needs to distinguish from the learner's trajectories.<jupyter_code>from imitation.data import rollout
from imitation.data.wrappers import RolloutInfoWrapper
from stable_baselines3.common.vec_env import DummyVecEnv

rollouts = rollout.rollout(
    expert,
    DummyVecEnv([lambda: RolloutInfoWrapper(gym.make("seals/CartPole-v0"))] * 5),
    rollout.make_sample_until(min_timesteps=None, min_episodes=60),
)<jupyter_output><empty_output><jupyter_text>Now we are ready to set up our AIRL trainer.Note, that the `reward_net` is actually the network of the discriminator.We evaluate the learner before and after training so we can see if it made any progress.<jupyter_code>from imitation.algorithms.adversarial.airl import AIRL
from imitation.rewards.reward_nets import BasicShapedRewardNet
from imitation.util.networks import RunningNorm
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv

import gym
import seals


venv = DummyVecEnv([lambda: gym.make("seals/CartPole-v0")] * 8)
learner = PPO(
    env=venv,
    policy=MlpPolicy,
    batch_size=64,
    ent_coef=0.0,
    learning_rate=0.0003,
    n_epochs=10,
)
reward_net = BasicShapedRewardNet(
    venv.observation_space, venv.action_space, normalize_input_layer=RunningNorm
)
airl_trainer = AIRL(
    demonstrations=rollouts,
    demo_batch_size=1024,
    gen_replay_buffer_capacity=2048,
    n_disc_updates_per_round=4,
    venv=venv,
    gen_algo=learner,
    reward_net=reward_net,
)

learner_rewards_before_training, _ = evaluate_policy(
    learner, venv, 100, return_episode_rewards=True
)
airl_trainer.train(20000)  # Note: set to 300000 for better results
learner_rewards_after_training, _ = evaluate_policy(
    learner, venv, 100, return_episode_rewards=True
)<jupyter_output><empty_output><jupyter_text>When we look at the histograms of rewards before and after learning, we can see that the learner is not perfect yet, but it made some progress at least.If not, just re-run the above cell.<jupyter_code>import matplotlib.pyplot as plt
import numpy as np

print(np.mean(learner_rewards_after_training))
print(np.mean(learner_rewards_before_training))

plt.hist(
    [learner_rewards_before_training, learner_rewards_after_training],
    label=["untrained", "trained"],
)
plt.legend()
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.07940904893813482, 'fraction_numerical': 0.02277623884272084, 'mean_word_length': 5.799163179916318, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '456627', 'license': 'MIT', 'n_tokens_mistral': 1160, 'n_tokens_neox': 1088, 'n_words': 270, 'path': 'examples/4_train_airl.ipynb', 'repo_name': 'mcx/imitation', 'size': '3249'}