huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

adding comments to gail_options_image, combining environments for options gail, and fixing bug where last state is yielded in hl buffer

Browse Source
This commit is contained in:
Arec
2021-10-20 12:28:12 -07:00
committed by Johannes Fischer
parent bef6d6db55
commit 03478b7ed9
2 changed files with 644 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

143
scratch/arec/intersimple/gail_options_image.py
View File

@@ -23,17 +23,38 @@ logging.basicConfig(level=logging.DEBUG)
ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
"""
Class for high-level options policy (generator)
"""
def __init__(self, observation_space, *args, **kwargs):
super().__init__(observation_space['obs'], *args, **kwargs)
def _prior_distribution(self, s):
latent_pi, latent_vf, latent_sde = self._get_latent(s)
distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
"""
Return prior distribution over high-level options (before masking)
Args:
s (torch.tensor): observation
Returns:
values (torch.tensor): values from critic
dist (torch.distributions): prior distribution over actions
"""
latent_pi, latent_vf, latent_sde = self._get_latent(s)
distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
values = self.value_net(latent_vf)
return values, distribution.distribution
def predict(self, obs):
"""
Will mask invalid states before making action selections
Args:
obs: dict with keys:
obs (torch.tensor): (B,o) true observations
mask (torch.tensor): (B,m) mask over valid actions
Returns:
ch (torch.tensor): (B,a) sampled actions
values (torch.tensor): (B,) predicted value at observation
log_probs (torch.tensor): (B,) log probabilities of selected actions
"""
s, m = obs['obs'], obs['mask']
values, prior = self._prior_distribution(s)
posterior = Categorical(prior.probs * m)
@@ -41,14 +62,31 @@ class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
return ch, values, posterior.log_prob(ch)
def evaluate_actions(self, obs, ch):
"""
Evaluate particular actions
Args:
obs: dict with keys:
obs (torch.tensor): (B,o) true observations
mask (torch.tensor): (B,m) masks over valid actions
ch (torch.tensor): (B,a) selected actions
Returns:
values (torch.tensor): (B,) predicted value at observation
log_probs (torch.tensor): (B,) log probabilities of selected actions
ent (torch.tensor): (B,) entropy of each distribution over actions
"""
s, m = obs['obs'], obs['mask']
values, prior = self._prior_distribution(s)
posterior = Categorical(prior.probs * m)
return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train
class OptionsEnv(gym.Wrapper):
"""
Wrap an intersimple environment with an options generator
"""
def __init__(self, env, *args, **kwargs):
"""
Initialize wrapped environment and set high-level action and observation spaces
"""
super().__init__(env, *args, **kwargs)
num_hl_options = len(ALL_OPTIONS)
self.action_space = gym.spaces.Discrete(num_hl_options)
@@ -67,51 +105,88 @@ class OptionsEnv(gym.Wrapper):
raise NotImplementedError('Use `LLOptions` or `HLOptions` for sampling.')
def sample(self, generator):
"""
yield transitions using a generator
Args:
generator (sb3.PPO)
Yields:
"""
self.done = True
while True:
self.episode_start = False
if self.done:
# reset environment
self.s = self.env.reset()
self.m = available_actions(self.env)
self.done = False
self.episode_start = True
# set the action, the value of the start state, and the logprob of the action
# according to the current environment state and mask
self.ch, self.value, self.log_prob = generator.policy.predict({
'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
})
# store a float list of actions to take given the option selected in the environment
self.plan = list(map(float, generate_plan(self.env, self.ch)))
# run whatever _after_choice might dictate in a child class
self._after_choice()
# some checks
assert not self.done
assert self.plan
assert feasible(self.env, self.plan, self.ch)
# execute the option so long as the episode isn't complete and the plan is still feasible
while not self.done and self.plan and feasible(self.env, self.plan, self.ch):
# pop first action
self.a, self.plan = self.plan[0], self.plan[1:]
# normalize action ??
self.a = self.env._normalize(self.a)
# step through environment
self.nexts, _, self.done, _ = self.env.step(self.a)
self.nextm = available_actions(self.env)
# run whatever _after_step might dictate in child class
self._after_step()
# update state and mask to current
self.s = self.nexts
self.m = self.nextm
# transitions yielded from self._transitions() functions specied in child classes
yield from self._transitions()
### NOTE: only yields after a full option has been executed / exited
class LLOptions(OptionsEnv):
"""Sample low-level (state, action) tuples for discriminator training."""
def __init__(self, *args, **kwargs):
"""
LLOption uses the true LL observations
"""
super().__init__(*args, **kwargs)
self.observation_space = self.observation_space['obs']
# overwrite observation space to just output obs directly
self.observation_space = self.observation_space['obs']
def _after_choice(self):
"""
After each option choice, initialize/reset the transition buffer
"""
self._transition_buffer = []
def _after_step(self):
"""
After each ll action, append s, s', a, done to transition buffer
"""
self._transition_buffer.append({
'obs': self.s,
'next_obs': self.nexts,
@@ -120,9 +195,17 @@ class LLOptions(OptionsEnv):
})
def _transitions(self):
"""
Yield from the transition buffer
"""
yield from self._transition_buffer
def sample_ll(self, policy):
"""
Not quite sure how this works????
Why would you do this over LLOptions.sample(policy)
"""
# What happens if you return a yield from ????????
return self.sample(policy)
class HLOptions(OptionsEnv):
@@ -132,10 +215,16 @@ class HLOptions(OptionsEnv):
super().__init__(*args, **kwargs)
def _after_choice(self):
"""
After an option selection, initialize total reward and number of steps
"""
self.r = 0
self.steps = 0
def _after_step(self):
"""
After each low-level action, add the discounted discriminated reward score (given a discriminator)
"""
self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train(
state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()),
action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()),
@@ -145,6 +234,18 @@ class HLOptions(OptionsEnv):
self.steps += 1
def _transitions(self):
"""
Yield a single dictionary per high-level selected action
Fields:
obs: high-level state and mask at selection
action: chosen high-level action
reward: accumulated option reward
episode_start: whether the action was chosen at the episode start
value: the value estimate from the starting state
log_prob: the log_prob of the selected action from the starting state
done: whether the episode has ended
"""
yield {
'obs': {'obs': self.s, 'mask': self.m},
'action': self.ch,
@@ -156,16 +257,29 @@ class HLOptions(OptionsEnv):
}
def sample_hl(self, policy, discriminator):
"""
Args:
policy
discriminator: function with which to score rewards
Returns:
gen: an which samples high-level transitions from the environment
"""
self.discriminator = discriminator
return self.sample(policy)
class RenderOptions(LLOptions):
def _after_step(self):
"""
Render the environment after each low-level step
"""
super()._after_step()
self.env.render()
def close(self, *args, **kwargs):
"""
On 'close', close the environment
"""
self.env.close(*args, **kwargs)
def available_actions(env):
@@ -307,6 +421,18 @@ def train_generator(env, generator, discriminator, num_samples):
generator.train()
def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99):
"""
Args:
expert_data: list of transitions
env_class: environment class
env_settings: environment settings
epochs: number of epochs to train for
discrim_batch_size: discriminator batch size
generator_steps: number of steps taken in generator
discount: discount factor
Returns:
generator (stable_baselines3.PPO): options policy
"""
env = env_class(**env_settings)
env.discount = discount
@@ -367,13 +493,12 @@ if __name__ == '__main__':
discount=0.99
)
generator.save(model_name)
generator.save(model_name) # save ppo sb3 generator class
# %%
model = stable_baselines3.PPO.load(model_name)
model = stable_baselines3.PPO.load(model_name) # not actually used
env = RenderOptions(NRasterizedRandomAgent(**env_settings))
for s in env.sample_ll(generator):
if s['dones']:
break