Sample expert data from same distribution

This commit is contained in:
ebuehrle
2021-11-04 18:52:08 +01:00
parent 9a95134530
commit f794a96112
4 changed files with 27 additions and 17 deletions

View File

@@ -16,31 +16,32 @@ from tqdm import tqdm
from src.policies.options import OptionsCnnPolicy
from src.gail.train import flatten_transitions
from gail.options2 import OptionsEnv, RenderOptions, imitation_discriminator
from gym.wrappers import TimeLimit
from gail.envs import TLNRasterizedRouteRandomAgentLocation
from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
model_name = 'gail_options_image_random_location'
env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True}
env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True, 'max_episode_steps': 50}
env = TLNRasterizedRouteRandomAgentLocation(**env_settings)
ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
def train(
expert_data,
expert_batch_size=4069,
expert_batch_size=2048,
discriminator_updates_per_round=10,
generator_steps=256,
generator_total_steps=2048,
generator_total_steps=1024,
generator_updates_per_round=10,
discount=0.99,
epochs=100,
):
env = NRasterizedRouteRandomAgentLocation(**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}/.")
venv = make_vec_env(NRasterizedRouteRandomAgentLocation, n_envs=1, env_kwargs=env_settings)
venv = DummyVecEnv([lambda: env])
discriminator = adversarial.GAIL(
expert_data=expert_data,
expert_batch_size=expert_batch_size,
@@ -50,13 +51,13 @@ def train(
gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
)
options_env = TimeLimit(OptionsEnv(
options_env = OptionsEnv(
env,
options=ALL_OPTIONS,
discriminator=imitation_discriminator(discriminator),
discount=discount,
ll_buffer_capacity=expert_batch_size,
), max_episode_steps=15)
)
generator = stable_baselines3.PPO(
OptionsCnnPolicy,
options_env,
@@ -70,9 +71,9 @@ def train(
generator.learn(total_timesteps=generator_total_steps)
# train discriminator
generator_samples = options_env.sample_ll(expert_batch_size)
generator_samples = flatten_transitions(generator_samples)
for _ in range(discriminator_updates_per_round):
generator_samples = options_env.sample_ll(expert_batch_size)
generator_samples = flatten_transitions(generator_samples)
discriminator.train_disc(gen_samples=generator_samples)
generator.save(model_name)
@@ -94,13 +95,12 @@ def video(model_name, env):
def evaluate():
video(
model_name=model_name,
env=NRasterizedRouteRandomAgentLocation(**env_settings)
env=env
)
# %%
if __name__ == '__main__':
with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001skip5.pkl", "rb") as f:
with open("data/NormalizedIntersimpleExpertMu.001N10000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
trajectories = pickle.load(f)
transitions = rollout.flatten_trajectories(trajectories)
train(transitions)