117 lines
4.0 KiB
Python
117 lines
4.0 KiB
Python
from gail.discriminator import MlpDiscriminator
|
|
from imitation.algorithms import adversarial
|
|
import stable_baselines3
|
|
import torch.utils.data
|
|
import numpy as np
|
|
from intersim.envs.intersimple import Intersimple
|
|
import itertools
|
|
from torch.distributions import Categorical
|
|
import gym
|
|
|
|
class OptionsMlpPolicy:
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self._policy = stable_baselines3.common.policies.ActorCriticPolicy(
|
|
*args, **kwargs
|
|
)
|
|
|
|
def _prior_distribution(self, s):
|
|
latent_pi, _, latent_sde = self._policy._get_latent(s)
|
|
distribution = self._policy._get_action_dist_from_latent(latent_pi, latent_sde)
|
|
return distribution.distribution
|
|
|
|
def predict(self, obs):
|
|
s, m = obs
|
|
prior = self._prior_distribution(s)
|
|
posterior = Categorical(prior.probs * m)
|
|
ch = posterior.sample()
|
|
return ch
|
|
|
|
def evaluate_actions(self, obs, ch):
|
|
s, m = obs
|
|
values = self._policy.value_net(s)
|
|
prior = self._prior_distribution(s)
|
|
posterior = Categorical(prior.probs * m)
|
|
return values, posterior.logprob(ch), posterior.entropy() # additional values used by PPO.train
|
|
|
|
def available_actions(env):
|
|
"""Return mask of available actions given current `env` state."""
|
|
return np.ones((env.num_hl_actions,))
|
|
|
|
def generate_plan(env, i):
|
|
"""Generate input profile for high-level action `i`."""
|
|
return np.zeros((env.num_hl_steps,))
|
|
|
|
def feasible(env, plan):
|
|
"""Check if input profile is feasible given current `env` state."""
|
|
return True
|
|
|
|
def sample_ll(env, generator):
|
|
"""Sample low-level (state, action) pairs for discriminator training."""
|
|
done = True
|
|
while True:
|
|
if done:
|
|
s = env.reset()
|
|
|
|
m = available_actions(env)
|
|
ch = generator.policy.predict((s, m))
|
|
plan = list(generate_plan(env, ch))
|
|
|
|
while not done and plan and feasible(env, plan):
|
|
a = plan.pop()
|
|
yield (s, a)
|
|
s, _, done, _ = env.step(a)
|
|
|
|
def train_discriminator(env, expert_data, generator, discriminator, generator_batch_size):
|
|
expert_samples = next(expert_data)
|
|
generator_samples = itertools.islice(sample_ll(env, generator), generator_batch_size)
|
|
discriminator.train_disc(expert_samples, generator_samples)
|
|
|
|
def sample_hl(env, generator, discriminator):
|
|
"""Sample high-level (state, action, reward) tuples for generator training."""
|
|
done = True
|
|
while True:
|
|
if done:
|
|
s = env.reset()
|
|
|
|
m = available_actions(env)
|
|
obs = (s, m)
|
|
ch = generator.policy.predict((s, m))
|
|
plan = list(generate_plan(env, ch))
|
|
r = 0
|
|
discount = 1
|
|
|
|
while not done and plan and feasible(env, plan):
|
|
a = plan.pop()
|
|
r += discount * discriminator.discrim_net(s, a)
|
|
discount *= env.discount
|
|
s, _, done, _ = env.step(a)
|
|
|
|
yield (obs, ch, r)
|
|
|
|
def train_generator(env, generator, discriminator, generator_batch_size):
|
|
generator_samples = itertools.islice(sample_hl(env, generator, discriminator), generator_batch_size)
|
|
generator.rollout_buffer.reset()
|
|
generator.rollout_buffer.add(generator_samples)
|
|
generator.train()
|
|
|
|
class OptionsEnv(gym.Wrapper):
|
|
|
|
def __init__(self, env):
|
|
super().__init__(env)
|
|
self.action_space = gym.spaces.Discrete(env.num_hl_options)
|
|
|
|
def train(expert_data, epochs=10, generator_batch_size=1024, expert_batch_size=1024, num_hl_options=10, num_hl_steps=10, discount=0.99):
|
|
env = Intersimple()
|
|
env.num_hl_options = num_hl_options
|
|
env.num_hl_steps = num_hl_steps
|
|
env.discount = discount
|
|
|
|
discriminator = adversarial.GAIL(discrim_kwargs={'discrim_net': MlpDiscriminator()})
|
|
generator = stable_baselines3.PPO(OptionsMlpPolicy, OptionsEnv(env))
|
|
expert_data = torch.utils.data.DataLoader(expert_data, expert_batch_size)
|
|
|
|
for _ in range(epochs):
|
|
train_discriminator(env, expert_data, generator, discriminator, generator_batch_size)
|
|
train_generator(env, generator, discriminator, generator_batch_size)
|