AdVIL tests

scratch/etienne/pillbox/learners/train.py

@@ -0,0 +1,248 @@
+from imitation.algorithms import adversarial, bc
+from imitation.util import logger, util
+from stable_baselines3 import PPO, DQN, SAC
+from soft_q import SQLPolicy
+from sqil import SQILReplayBuffer
+from stable_baselines3.common import policies
+from stable_baselines3.common.evaluation import evaluate_policy
+from imitation.rewards import discrim_nets
+import numpy as np
+import argparse
+from utils import make_sa_dataloader, make_sads_dataloader, make_sa_dataset, linear_schedule
+from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
+from adril import AdRILWrapper, AdRILReplayBuffer
+import os
+from gym.spaces import Discrete
+import gym
+from advil import advil_training
+from stable_baselines3.common.running_mean_std import RunningMeanStd
+
+from advil import AdVILPolicy, AdVILDiscriminator
+
+def train_bc(env, n=0):
+    venv = util.make_vec_env(env, n_envs=8)
+    if isinstance(venv.action_space, Discrete):
+        w = 64
+    else:
+        w = 256
+    for i in range(n):
+        mean_rewards = []
+        std_rewards = []
+        for num_trajs in range(0, 26, 5):
+            if num_trajs == 0:
+                expert_data = make_sa_dataloader(env, normalize=False)
+            else:
+                expert_data = make_sa_dataloader(env, max_trajs=num_trajs, normalize=False)
+            bc_trainer = bc.BC(venv.observation_space, venv.action_space, expert_data=expert_data,
+                               policy_class=policies.ActorCriticPolicy,
+                               ent_weight=0., l2_weight=0., policy_kwargs=dict(net_arch=[w, w]))
+            if num_trajs > 0:
+                bc_trainer.train(n_batches=int(5e5))
+
+            def get_policy(*args, **kwargs):
+                return bc_trainer.policy
+            model = PPO(get_policy, env, verbose=1)
+            model.save(os.path.join("learners", env,
+                                    "bc_{0}_{1}".format(i, num_trajs)))
+            mean_reward, std_reward = evaluate_policy(
+                    model, model.get_env(), n_eval_episodes=10)               
+            mean_rewards.append(mean_reward)
+            std_rewards.append(std_reward)
+            print("{0} Trajs: {1}".format(num_trajs, mean_reward))
+            np.savez(os.path.join("learners", env, "bc_rewards_{0}".format(
+                i)), means=mean_rewards, stds=std_rewards)
+
+
+def train_gail(env, n=0):
+    venv = util.make_vec_env(env, n_envs=8)
+    if isinstance(venv.action_space, Discrete):
+        w = 64
+    else:
+        w = 256
+    expert_data = make_sads_dataloader(env, max_trajs=5)
+    logger.configure(os.path.join("learners", "GAIL"))
+
+    for i in range(n):
+        discrim_net = discrim_nets.ActObsMLP(
+                action_space=venv.action_space,
+                observation_space=venv.observation_space,
+                hid_sizes=(w, w),
+                )
+        gail_trainer = adversarial.GAIL(venv, expert_data=expert_data, expert_batch_size=32,
+                                        gen_algo=PPO("MlpPolicy", venv, verbose=1, n_steps=1024,
+                                                     policy_kwargs=dict(net_arch=[w, w])),
+                                                     discrim_kwargs={'discrim_net': discrim_net})
+        mean_rewards = []
+        std_rewards = []
+        for train_steps in range(20):
+            if train_steps > 0:
+                if 'Bullet' in env:
+                    gail_trainer.train(total_timesteps=25000)
+                else:
+                    gail_trainer.train(total_timesteps=16384)
+
+            def get_policy(*args, **kwargs):
+                return gail_trainer.gen_algo.policy
+            model = PPO(get_policy, env, verbose=1)
+            mean_reward, std_reward = evaluate_policy(
+                model, model.env, n_eval_episodes=10)
+            mean_rewards.append(mean_reward)
+            std_rewards.append(std_reward)
+            print("{0} Steps: {1}".format(train_steps, mean_reward))
+            np.savez(os.path.join("learners", env, "gail_rewards_{0}".format(i)),
+                     means=mean_rewards, stds=std_rewards)
+
+
+def train_sqil(env, n=0):
+    venv = gym.make(env)
+    expert_data = make_sa_dataset(env, max_trajs=5)
+
+    for i in range(n):
+        if isinstance(venv.action_space, Discrete):
+            model = DQN(SQLPolicy, venv, verbose=1, policy_kwargs=dict(net_arch=[64, 64]), learning_starts=1)
+        else:
+            model = SAC('MlpPolicy', venv, verbose=1, policy_kwargs=dict(net_arch=[256, 256]), ent_coef='auto',
+                        learning_rate=linear_schedule(7.3e-4), train_freq=64, gradient_steps=64, gamma=0.98, tau=0.02)
+
+        model.replay_buffer = SQILReplayBuffer(model.buffer_size, model.observation_space,
+                                               model.action_space, model.device, 1,
+                                               model.optimize_memory_usage, expert_data=expert_data)
+        mean_rewards = []
+        std_rewards = []
+        for train_steps in range(20):
+            if train_steps > 0:
+                if 'Bullet' in env:
+                    model.learn(total_timesteps=25000, log_interval=1)
+                else:
+                    model.learn(total_timesteps=16384, log_interval=1)
+            mean_reward, std_reward = evaluate_policy(
+                model, model.env, n_eval_episodes=10)
+            mean_rewards.append(mean_reward)
+            std_rewards.append(std_reward)
+            print("{0} Steps: {1}".format(train_steps, mean_reward))
+            np.savez(os.path.join("learners", env, "sqil_rewards_{0}".format(i)),
+                     means=mean_rewards, stds=std_rewards)
+
+
+def train_adril(env, n=0, balanced=False):
+    num_trajs = 20
+    expert_data = make_sa_dataset(env, max_trajs=num_trajs)
+    n_expert = len(expert_data["obs"])
+    expert_sa = np.concatenate((expert_data["obs"], np.reshape(expert_data["acts"], (n_expert, -1))), axis=1)
+
+    for i in range(0, n):
+        venv = AdRILWrapper(gym.make(env))
+        mean_rewards = []
+        std_rewards = []
+        # Create model
+        if isinstance(venv.action_space, Discrete):
+            model = DQN(SQLPolicy, venv, verbose=1, policy_kwargs=dict(net_arch=[64, 64]), learning_starts=1)
+        else:
+            model = SAC('MlpPolicy', venv, verbose=1, policy_kwargs=dict(net_arch=[256, 256]), ent_coef='auto',
+                        learning_rate=linear_schedule(7.3e-4), train_freq=64, gradient_steps=64, gamma=0.98, tau=0.02)
+        model.replay_buffer = AdRILReplayBuffer(model.buffer_size, model.observation_space,
+                                               model.action_space, model.device, 1,
+                                               model.optimize_memory_usage, expert_data=expert_data, N_expert=num_trajs,
+                                               balanced=balanced)
+        if not balanced:
+            for j in range(len(expert_sa)):
+                obs = expert_data["obs"][j]
+                act = expert_data["acts"][j]
+                next_obs = expert_data["next_obs"][j]
+                done = expert_data["dones"][j]
+                model.replay_buffer.add(obs, next_obs, act, -1, done)    
+        for train_steps in range(400):
+            # Train policy
+            if train_steps > 0:
+                if 'Bullet' in env:
+                    model.learn(total_timesteps=1250, log_interval=1000)
+                else:
+                    model.learn(total_timesteps=25000, log_interval=1000)
+                if train_steps % 1 == 0: # written to support more complex update schemes
+                    model.replay_buffer.set_iter(train_steps)
+                    model.replay_buffer.set_n_learner(venv.num_trajs)
+            
+            # Evaluate policy
+            if train_steps % 20 == 0:
+                model.set_env(gym.make(env))
+                mean_reward, std_reward = evaluate_policy(
+                    model, model.env, n_eval_episodes=10)
+                mean_rewards.append(mean_reward)
+                std_rewards.append(std_reward)
+                print("{0} Steps: {1}".format(int(train_steps * 1250), mean_reward))
+                np.savez(os.path.join("learners", env, "adril_rewards_{0}".format(i)),
+                        means=mean_rewards, stds=std_rewards)
+            # Update env
+            if train_steps > 0:
+                if train_steps % 1  == 0:
+                    venv.set_iter(train_steps + 1)
+            model.set_env(venv)
+
+
+def train_advil(env, policy_class=AdVILPolicy, discriminator_class=AdVILDiscriminator,
+                iters=int(1e5), lr_pi=8e-6, lr_f=8e-4):
+    venv = gym.make(env)
+    expert_data = make_sa_dataloader(
+        env,
+        normalize=False,
+        batch_size=1024,
+    )
+    pi = advil_training(
+        expert_data,
+        venv,
+        iters=iters,
+        policy_class=policy_class,
+        discriminator_class=discriminator_class, 
+        lr_pi=lr_pi,
+        lr_f=lr_f,
+    )
+    return pi
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Train expert policies.')
+    parser.add_argument(
+        '-a', '--algo', choices=['bc', 'gail', 'sqil', 'adril', 'advil', 'all'], required=True)
+    parser.add_argument('-e', '--env', choices=['cartpole', 'lunarlander', 'acrobot', 'pendulum', 'halfcheetah', 'walker', 'hopper', 'ant'],
+                        required=True)
+    parser.add_argument('-n', '--num_runs', required=False)
+    args = parser.parse_args()
+    if args.env == "cartpole":
+        envname = 'CartPole-v1'
+    elif args.env == "lunarlander":
+        envname = 'LunarLander-v2'
+    elif args.env == "acrobot":
+        envname = 'Acrobot-v1'
+    elif args.env == "pendulum":
+        envname = 'Pendulum-v0'
+    elif args.env == "halfcheetah":
+        envname = 'HalfCheetahBulletEnv-v0'
+    elif args.env == "walker":
+        envname = 'Walker2DBulletEnv-v0'
+    elif args.env == "hopper":
+        envname = 'HopperBulletEnv-v0'
+    elif args.env == "ant":
+        envname = 'AntBulletEnv-v0'
+    else:
+        print("ERROR: unsupported env.")
+    if args.num_runs is not None and args.num_runs.isdigit():
+        num_runs = int(args.num_runs)
+    else:
+        num_runs = 1
+    if args.algo == 'bc':
+        train_bc(envname, num_runs)
+    elif args.algo == 'gail':
+        train_gail(envname, num_runs)
+    elif args.algo == 'sqil':
+        train_sqil(envname, num_runs)
+    elif args.algo == 'adril':
+        train_adril(envname, num_runs)
+    elif args.algo == 'advil':
+        train_advil(envname, num_runs)
+    elif args.algo == 'all':
+        train_bc(envname, num_runs)
+        train_gail(envname, num_runs)
+        train_sqil(envname, num_runs)
+        train_adril(envname, num_runs)
+        train_advil(envname, num_runs)
+    else:
+        print("ERROR: unsupported algorithm")