<jupyter_start><jupyter_text>Learning a Reward Function using Kernel DensityThis demo shows how to train a `Pendulum` agent (exciting!) with our simple density-based imitation learning baselines. `DensityTrainer` has a few interesting parameters, but the key ones are:1. `density_type`: this governs whether density is measured on $(s,s')$ pairs (`db.STATE_STATE_DENSITY`), $(s,a)$ pairs (`db.STATE_ACTION_DENSITY`), or single states (`db.STATE_DENSITY`).2. `is_stationary`: determines whether a separate density model is used for each time step $t$ (`False`), or the same model is used for transitions at all times (`True`).3. `standardise_inputs`: if `True`, each dimension of the agent state vectors will be normalised to have zero mean and unit variance over the training dataset. This can be useful when not all elements of the demonstration vector are on the same scale, or when some elements have too wide a variation to be captured by the fixed kernel width (1 for Gaussian kernel).4. `kernel`: changes the kernel used for non-parametric density estimation. `gaussian` and `exponential` are the best bets; see the [sklearn docs](https://scikit-learn.org/stable/modules/density.htmlkernel-density) for the rest.<jupyter_code>import pprint

from imitation.algorithms import density as db
from imitation.data import types
from imitation.util import util
# Set FAST = False for longer training. Use True for testing and CI.
FAST = True

if FAST:
    N_VEC = 1
    N_TRAJECTORIES = 1
    N_ITERATIONS = 1
    N_RL_TRAIN_STEPS = 100

else:
    N_VEC = 8
    N_TRAJECTORIES = 10
    N_ITERATIONS = 100
    N_RL_TRAIN_STEPS = int(1e5)
from stable_baselines3.common.policies import ActorCriticPolicy
from stable_baselines3 import PPO

env_name = "Pendulum-v1"
env = util.make_vec_env(env_name, N_VEC)
rollouts = types.load("../tests/testdata/expert_models/pendulum_0/rollouts/final.pkl")


imitation_trainer = PPO(ActorCriticPolicy, env, learning_rate=3e-4, n_steps=2048)
density_trainer = db.DensityAlgorithm(
    venv=env,
    demonstrations=rollouts,
    rl_algo=imitation_trainer,
    density_type=db.DensityType.STATE_ACTION_DENSITY,
    is_stationary=True,
    kernel="gaussian",
    kernel_bandwidth=0.2,  # found using divination & some palm reading
    standardise_inputs=True,
)
density_trainer.train()
novice_stats = density_trainer.test_policy()
print("Novice stats (true reward function):")
pprint.pprint(novice_stats)
novice_stats_im = density_trainer.test_policy(
    true_reward=False, n_trajectories=N_TRAJECTORIES
)
print("Novice stats (imitation reward function):")
pprint.pprint(novice_stats_im)

for i in range(N_ITERATIONS):
    density_trainer.train_policy(N_RL_TRAIN_STEPS)

    good_stats = density_trainer.test_policy(n_trajectories=N_TRAJECTORIES)
    print(f"Trained stats (epoch {i}):")
    pprint.pprint(good_stats)
    novice_stats_im = density_trainer.test_policy(true_reward=False)
    print(f"Trained stats (imitation reward function, epoch {i}):")
    pprint.pprint(novice_stats_im)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08699934340118187, 'fraction_numerical': 0.010177281680892974, 'mean_word_length': 5.387840670859539, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '78413', 'license': 'MIT', 'n_tokens_mistral': 1033, 'n_tokens_neox': 985, 'n_words': 289, 'path': 'examples/7_train_density.ipynb', 'repo_name': 'mcx/imitation', 'size': '3046'}