from stable_baselines3 import PPO
from stable_baselines3.ppo import MlpPolicy

import gym

env = gym.make("CartPole-v1")
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

"""
Explanation: Train an Agent using the DAgger Algorithm
The DAgger algorithm is an extension of behavior cloning. 
In behavior cloning, the training trajectories are recorded directly from an expert.
In DAgger, the learner generates the trajectories but an experts corrects the actions with the optimal actions in each of the visited states.
This ensures that the state distribution of the training data matches that of the learner's current policy.
First we need an expert to learn from:
End of explanation
"""


import tempfile
import gym
from stable_baselines3.common.vec_env import DummyVecEnv

from imitation.algorithms import bc
from imitation.algorithms.dagger import SimpleDAggerTrainer

venv = DummyVecEnv([lambda: gym.make("CartPole-v1")])


bc_trainer = bc.BC(
    observation_space=env.observation_space,
    action_space=env.action_space,
)

with tempfile.TemporaryDirectory(prefix="dagger_example_") as tmpdir:
    print(tmpdir)
    dagger_trainer = SimpleDAggerTrainer(
        venv=venv, scratch_dir=tmpdir, expert_policy=expert, bc_trainer=bc_trainer
    )

    dagger_trainer.train(2000)

"""
Explanation: Then we can construct a DAgger trainer und use it to train the policy on the cartpole environment.
End of explanation
"""


from stable_baselines3.common.evaluation import evaluate_policy

reward, _ = evaluate_policy(dagger_trainer.policy, env, 10)
print(reward)

"""
Explanation: Finally, the evaluation shows, that we actually trained a policy that solves the environment (500 is the max reward).
End of explanation
"""


