Copy files

src/core/gail.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn.functional as F
+from dataclasses import dataclass
+from core.reparam_module import ReparamPolicy
+from core.sampling import rollout
+from core.trpo import trpo_step
+from core.ppo import ppo_step
+from tqdm import tqdm
+
+class TerminalLogger:
+    def add_scalar(self, key, scalar, i=None):
+        if i is not None:
+            print('Iteration', i, end=' ')
+        print(key, scalar)
+
+@dataclass
+class Buffer:
+    states: torch.Tensor
+    actions: torch.Tensor
+    rewards: torch.Tensor
+    dones: torch.Tensor
+
+def roll_buffer(buffer, *args, **kwargs):
+    return Buffer(
+        torch.roll(buffer.states, *args, **kwargs),
+        torch.roll(buffer.actions, *args, **kwargs),
+        torch.roll(buffer.rewards, *args, **kwargs),
+        torch.roll(buffer.dones, *args, **kwargs),
+    )
+
+def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value,
+         v_opt, v_iters, epochs, rollout_episodes, rollout_steps, gamma,
+         gae_lambda, delta, backtrack_coeff, backtrack_iters, cg_iters=10, cg_damping=0.1, wasserstein=False, wasserstein_c=None, logger=TerminalLogger()):
+
+    policy(torch.zeros(env_fn(0).observation_space.shape))
+    policy = ReparamPolicy(policy)
+
+    logger.add_scalar('expert/mean_episode_length', (~expert_data.dones).sum() / expert_data.states.shape[0])
+    logger.add_scalar('expert/mean_reward_per_episode', expert_data.rewards[~expert_data.dones].sum() / expert_data.states.shape[0])
+
+    for epoch in tqdm(range(epochs)):
+        generator_data = Buffer(*rollout(env_fn, policy, rollout_episodes, rollout_steps))
+
+        logger.add_scalar('gen/mean_episode_length', (~generator_data.dones).sum() / generator_data.states.shape[0], epoch)
+        logger.add_scalar('gen/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+
+        discriminator, loss = train_discriminator(expert_data, generator_data, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c)
+        if wasserstein:
+            generator_data.rewards = discriminator(generator_data.states, generator_data.actions)
+        else:
+            generator_data.rewards = -F.logsigmoid(discriminator(generator_data.states, generator_data.actions))
+        logger.add_scalar('disc/final_loss', loss, epoch)
+        logger.add_scalar('disc/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+
+        value, policy = trpo_step(value, policy, generator_data.states, generator_data.actions, generator_data.rewards, generator_data.dones, gamma, gae_lambda, delta, backtrack_coeff, backtrack_iters, v_opt, v_iters, cg_iters, cg_damping)
+        expert_data = roll_buffer(expert_data, shifts=-3, dims=0)
+    
+    return value, policy
+
+def gail_ppo(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value,
+         v_opt, v_iters, epochs, rollout_episodes, rollout_steps, gamma,
+         gae_lambda, clip_ratio, pi_opt, pi_iters, target_kl=None, max_grad_norm=None, wasserstein=False, wasserstein_c=None, logger=TerminalLogger()):
+
+    logger.add_scalar('expert/mean_episode_length', (~expert_data.dones).sum() / expert_data.states.shape[0])
+    logger.add_scalar('expert/mean_reward_per_episode', expert_data.rewards[~expert_data.dones].sum() / expert_data.states.shape[0])
+
+    for epoch in range(epochs):
+        generator_data = Buffer(*rollout(env_fn, policy, rollout_episodes, rollout_steps))
+
+        logger.add_scalar('gen/mean_episode_length', (~generator_data.dones).sum() / generator_data.states.shape[0], epoch)
+        logger.add_scalar('gen/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+
+        discriminator, loss = train_discriminator(expert_data, generator_data, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c)
+        if wasserstein:
+            generator_data.rewards = discriminator(generator_data.states, generator_data.actions)
+        else:
+            generator_data.rewards = -F.logsigmoid(discriminator(generator_data.states, generator_data.actions))
+        logger.add_scalar('disc/final_loss', loss, epoch)
+        logger.add_scalar('disc/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+
+        value, policy = ppo_step(value, policy, generator_data.states, generator_data.actions, generator_data.rewards, generator_data.dones, clip_ratio, gamma, gae_lambda, pi_opt, pi_iters, v_opt, v_iters, target_kl, max_grad_norm)
+        expert_data = roll_buffer(expert_data, shifts=-3, dims=0)
+    
+    return value, policy
+
+def train_discriminator(expert_data, generator_data, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c=None):
+
+    n_expert_samples = (~expert_data.dones).sum()
+    n_generator_samples = (~generator_data.dones).sum()
+    n_samples = torch.minimum(n_expert_samples, n_generator_samples)
+
+    gen_states = generator_data.states[~generator_data.dones][:n_samples]
+    gen_actions = generator_data.actions[~generator_data.dones][:n_samples]
+    exp_states = expert_data.states[~expert_data.dones][:n_samples]
+    exp_actions = expert_data.actions[~expert_data.dones][:n_samples]
+
+    states = torch.cat((exp_states, gen_states), dim=0).detach()
+    actions = torch.cat((exp_actions, gen_actions), dim=0).detach()
+    labels = torch.cat((torch.zeros(n_samples), torch.ones(n_samples))).detach()
+
+    # print('Batch augmentation on')
+    # random_states = torch.rand_like(gen_states)
+    # random_actions = torch.rand_like(gen_actions)
+    # states = torch.cat((exp_states, gen_states, random_states), dim=0).detach()
+    # actions = torch.cat((exp_actions, gen_actions, random_actions), dim=0).detach()
+    # labels = torch.cat((torch.zeros(n_samples), torch.ones(n_samples), torch.ones(n_samples))).detach()
+
+    for _ in range(disc_iters):
+        disc_opt.zero_grad()
+        pred = discriminator(states, actions)
+
+        if wasserstein:
+            loss = -(pred * (1 - labels) - pred * labels).mean()
+        else:
+            loss = F.binary_cross_entropy(torch.sigmoid(pred), labels)
+        
+        loss.backward()
+        disc_opt.step()
+
+        if wasserstein_c is not None:
+            with torch.no_grad():
+                for param in discriminator.parameters():
+                    param.clamp_(-wasserstein_c, wasserstein_c)
+    
+    return discriminator, loss