72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
# %%
|
|
import sys
|
|
sys.path.append('../../../')
|
|
|
|
import pickle
|
|
import imitation.data.rollout as rollout
|
|
import imitation.data.types as types
|
|
import torch
|
|
from gail.envs import TLNRasterizedRouteRandomAgentLocation
|
|
import tempfile
|
|
import pathlib
|
|
from imitation.util import logger
|
|
from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
|
|
from imitation.algorithms import adversarial
|
|
from src.discriminator import CnnDiscriminator
|
|
import stable_baselines3
|
|
from tqdm import tqdm
|
|
|
|
with open("data/NormalizedIntersimpleExpertMu.001N50000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
|
|
trajectories = pickle.load(f)
|
|
transitions = rollout.flatten_trajectories(trajectories)
|
|
|
|
# %%
|
|
env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True, 'max_episode_steps': 200}
|
|
env = TLNRasterizedRouteRandomAgentLocation(**env_settings)
|
|
|
|
tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
|
|
tempdir_path = pathlib.Path(tempdir.name)
|
|
logger.configure(tempdir_path / "GAIL/")
|
|
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
|
|
|
|
expert_batch_size = 4096
|
|
|
|
venv = DummyVecEnv([lambda: env])
|
|
discriminator = adversarial.GAIL(
|
|
expert_data=transitions,
|
|
expert_batch_size=expert_batch_size,
|
|
#discrim_kwargs={'discrim_net': NoisyDiscriminator(venv, std=0.25)},
|
|
disc_opt_cls=torch.optim.RMSprop,
|
|
disc_opt_kwargs={'lr': 0.0001, 'weight_decay': 0.003},
|
|
discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
|
|
venv=venv, # unused
|
|
gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
|
|
)
|
|
|
|
expert_data_loader = torch.utils.data.DataLoader(
|
|
transitions,
|
|
batch_size=expert_batch_size,
|
|
collate_fn=types.transitions_collate_fn,
|
|
shuffle=True,
|
|
drop_last=True,
|
|
)
|
|
|
|
gen_data_loader = torch.utils.data.DataLoader(
|
|
transitions,
|
|
batch_size=expert_batch_size,
|
|
collate_fn=types.transitions_collate_fn,
|
|
shuffle=True,
|
|
drop_last=True,
|
|
)
|
|
|
|
# %%
|
|
epochs = 1000
|
|
for i in tqdm(range(epochs)):
|
|
for expert_samples, gen_samples in zip(expert_data_loader, gen_data_loader):
|
|
# randomly corrupt actions
|
|
gen_samples['acts'] = -1 + 2 * torch.rand(*gen_samples['acts'].shape)
|
|
|
|
discriminator.train_disc(expert_samples=expert_samples, gen_samples=gen_samples)
|
|
|
|
torch.save(discriminator.discrim_net.state_dict(), 'train_discrim.pt')
|