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

Merge branch 'main' of https://github.com/sisl/InteractionImitation

Browse Source
This commit is contained in:
ebuehrle
2022-02-26 14:06:22 +01:00
parent 6876cf9625 c28c6c05b7
commit 1e24612347
16 changed files with 289 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
src/core/policy.py
View File

@@ -36,30 +36,31 @@ class BasePolicy(nn.Module):
class Policy(BasePolicy):
def __init__(self, *args, **kwargs):
def __init__(self, *args, hidden_layer_size=50, n_hidden_layers=2, activation=nn.Tanh, **kwargs):
super().__init__(*args, **kwargs)
self.nn = nn.Sequential(
nn.LazyLinear(50),
nn.Tanh(),
nn.LazyLinear(50),
nn.Tanh(),
nn.LazyLinear(2 * self.action_dim),
)
layers = sum([[nn.LazyLinear(hidden_layer_size),
activation()] for _ in range(n_hidden_layers)],[])
self.nn = nn.Sequential(*layers, nn.LazyLinear(2 *self.action_dim))
# old
# self.nn = nn.Sequential(
# nn.LazyLinear(50),
# nn.Tanh(),
# nn.LazyLinear(50),
# nn.Tanh(),
# nn.LazyLinear(2 * self.action_dim),
#)
def forward(self, states):
return self.nn(states)
class DiscretePolicy(BasePolicy):
def __init__(self, *args, hidden_layer_size=50, **kwargs):
def __init__(self, *args, hidden_layer_size=50, n_hidden_layers=2, activation=nn.Tanh, **kwargs):
super().__init__(*args, **kwargs)
self.nn = nn.Sequential(
nn.LazyLinear(hidden_layer_size),
nn.Tanh(),
nn.LazyLinear(hidden_layer_size),
nn.Tanh(),
nn.LazyLinear(self.action_dim),
)
layers = sum([[nn.LazyLinear(hidden_layer_size),
activation()] for _ in range(n_hidden_layers)],[])
self.nn = nn.Sequential(*layers, nn.LazyLinear(self.action_dim))
def forward(self, states):
return self.nn(states)

17
src/eval_main.py
View File

@@ -37,21 +37,22 @@ def load_policy(method:str,
Returns:
policy (Optional[BaseAlgorithm]): the policy to evaluate
"""
ml = torch.device('cpu') if not torch.cuda.is_available() else None
if method == 'idm':
policy = IDMRulePolicy(env, **policy_kwargs)
elif method == 'bc':
policy = SetPolicy(env.action_space.shape[-1])
policy.load_state_dict(torch.load(policy_file))
policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval()
elif method == 'gail':
policy = SetPolicy(env.action_space.shape[-1])
policy(torch.zeros(env.observation_space.shape))
policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load(policy_file))
policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval()
elif method == 'gail-ppo':
policy = SetPolicy(env.action_space.shape[-1])
policy.load_state_dict(torch.load(policy_file))
policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval()
elif method == 'rail':
raise NotImplementedError
@@ -59,11 +60,11 @@ def load_policy(method:str,
policy = SetDiscretePolicy(env.action_space.n)
policy(torch.zeros(env.observation_space.shape))
policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load(policy_file))
policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval()
elif method == 'ogail-ppo':
policy = SetDiscretePolicy(env.action_space.n)
policy.load_state_dict(torch.load(policy_file))
policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval()
elif method == 'sgail':
policy = SetMaskedDiscretePolicy(env.action_space.n)
@@ -72,11 +73,11 @@ def load_policy(method:str,
torch.zeros(env.observation_space['safe_actions'].shape)
)
policy = ReparamSafePolicy(policy)
policy.load_state_dict(torch.load(policy_file))
policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval()
elif method == 'sgail-ppo':
policy = SetMaskedDiscretePolicy(env.action_space.n)
policy.load_state_dict(torch.load(policy_file))
policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval()
else:
raise NotImplementedError
@@ -406,4 +407,4 @@ def eval_main(
if __name__=='__main__':
import fire
fire.Fire(eval_main)
fire.Fire(eval_main)

27
src/safe_options/options.py
View File

@@ -52,7 +52,8 @@ def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value
generator_data.ll.actions += 0.1 * torch.randn_like(generator_data.ll.actions)
logger.add_scalar('gen/mean_episode_length', (~generator_data.ll.dones).sum() / generator_data.ll.states.shape[0], epoch)
gen_mean_episode_length = (~generator_data.ll.dones).sum() / generator_data.ll.states.shape[0]
logger.add_scalar('gen/mean_episode_length', gen_mean_episode_length , epoch)
gen_mean_reward_per_episode = generator_data.hl.rewards[~generator_data.hl.dones].sum() / generator_data.hl.states.shape[0]
logger.add_scalar('gen/mean_reward_per_episode', gen_mean_reward_per_episode, epoch)
logger.add_scalar('gen/unsafe_probability_mass', policy.unsafe_probability_mass(policy(generator_data.hl.states[~generator_data.hl.dones], generator_data.hl.safe_actions[~generator_data.hl.dones])).mean(), epoch)
@@ -64,7 +65,8 @@ def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value
else:
generator_data.ll.rewards = -F.logsigmoid(discriminator(generator_data.ll.states, generator_data.ll.actions))
logger.add_scalar('disc/final_loss', loss, epoch)
logger.add_scalar('disc/mean_reward_per_episode', generator_data.ll.rewards[~generator_data.ll.dones].sum() / generator_data.ll.states.shape[0], epoch)
disc_mean_reward_per_episode = generator_data.ll.rewards[~generator_data.ll.dones].sum() / generator_data.ll.states.shape[0]
logger.add_scalar('disc/mean_reward_per_episode', disc_mean_reward_per_episode , epoch)
#assert generator_data.ll.rewards.shape == generator_data.ll.dones.shape
generator_data.hl.rewards = torch.where(~generator_data.ll.dones, generator_data.ll.rewards, torch.tensor(0.)).sum(-1)
@@ -77,7 +79,9 @@ def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value
'epoch': epoch,
'value': value,
'policy': policy,
'gen/mean_reward_per_episode': gen_mean_reward_per_episode,
'gen/mean_episode_length': gen_mean_episode_length.item(),
'gen/mean_reward_per_episode': gen_mean_reward_per_episode.item(),
'disc/mean_reward_per_episode': disc_mean_reward_per_episode.item(),
})
return value, policy
@@ -94,8 +98,8 @@ def gail_ppo(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, v
generator_data = OptionsRollout(HLBuffer(*hl_data), Buffer(*ll_data))
generator_data.ll.actions += 0.1 * torch.randn_like(generator_data.ll.actions)
logger.add_scalar('gen/mean_episode_length', (~generator_data.ll.dones).sum() / generator_data.ll.states.shape[0], epoch)
gen_mean_episode_length = (~generator_data.ll.dones).sum() / generator_data.ll.states.shape[0]
logger.add_scalar('gen/mean_episode_length', gen_mean_episode_length, epoch)
gen_mean_reward_per_episode = generator_data.hl.rewards[~generator_data.hl.dones].sum() / generator_data.hl.states.shape[0]
logger.add_scalar('gen/mean_reward_per_episode', gen_mean_reward_per_episode, epoch)
logger.add_scalar('gen/unsafe_probability_mass', policy.unsafe_probability_mass(policy(generator_data.hl.states[~generator_data.hl.dones], generator_data.hl.safe_actions[~generator_data.hl.dones])).mean(), epoch)
@@ -107,7 +111,8 @@ def gail_ppo(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, v
else:
generator_data.ll.rewards = -F.logsigmoid(discriminator(generator_data.ll.states, generator_data.ll.actions))
logger.add_scalar('disc/final_loss', loss, epoch)
logger.add_scalar('disc/mean_reward_per_episode', generator_data.ll.rewards[~generator_data.ll.dones].sum() / generator_data.ll.states.shape[0], epoch)
disc_mean_reward_per_episode = generator_data.ll.rewards[~generator_data.ll.dones].sum() / generator_data.ll.states.shape[0]
logger.add_scalar('disc/mean_reward_per_episode', disc_mean_reward_per_episode, epoch)
#assert generator_data.ll.rewards.shape == generator_data.ll.dones.shape
generator_data.hl.rewards = torch.where(~generator_data.ll.dones, generator_data.ll.rewards, torch.tensor(0.)).sum(-1)
@@ -120,7 +125,9 @@ def gail_ppo(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, v
'epoch': epoch,
'value': value,
'policy': policy,
'gen/mean_reward_per_episode': gen_mean_reward_per_episode,
'gen/mean_episode_length': gen_mean_episode_length.item(),
'gen/mean_reward_per_episode': gen_mean_reward_per_episode.item(),
'disc/mean_reward_per_episode': disc_mean_reward_per_episode.item(),
})
for lr_scheduler in lr_schedulers:
@@ -248,9 +255,9 @@ class SafeOptionsEnv(OptionsEnv):
if d:
break
if self.abort_unsafe_collision_method is not None and \
not feasible(self.env, plan[k:], method=self.abort_unsafe_collision_method):
break
if self.abort_unsafe_collision_method is not None:
if not feasible(self.env, plan[k:], method=self.abort_unsafe_collision_method):
break
n_steps = k + 1
return observations, actions, rewards, env_done, plan_done, infos, n_steps