first commit

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Ankesh Anand
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,35 @@
+# State Representation Learning Using an Unbalanced Atlas
+
+The project is based on the code from the benchmark and techniques introduced in the paper [Unsupervised State Representation Learning in Atari](https://arxiv.org/abs/1906.08226). Please visit https://github.com/mila-iqia/atari-representation-learning for detailed instructions on the benchmark.
+
+To run the script:
+
+```bash
+python run_probe.py
+```
+
+An example of running DIM-UA and setting the environment to Video Pinball, 4 heads and 512 units each, seed 2:
+
+```bash
+python run_probe.py --env-name VideoPinballNoFrameskip-v4 --n-head 4 --feature-size 512 --qv --seed 2
+```
+
+An example of running ST-DIM and setting the environment to Video Pinball, 512 units, seed 2:
+
+```bash
+python run_probe.py --env-name VideoPinballNoFrameskip-v4 --n-head 1 --feature-size 512 --seed 2
+```
+
+Running '-UA' described in the paper, and setting the environment to Video Pinball, 4 heads and 512 units each, seed 2:
+
+```bash
+python run_probe.py --env-name VideoPinballNoFrameskip-v4 --n-head 4 --feature-size 512 --seed 2
+```
+
+Running '+MMD' described in the paper, and setting the environment to Video Pinball, 4 heads and 512 units each, seed 2:
+
+```bash
+python run_probe.py --env-name VideoPinballNoFrameskip-v4 --n-head 4 --feature-size 512 --mmd --seed 2
+```
+
+A detailed list of parameter setup is in [atariari/methods/utils.py]

a2c_ppo_acktr/algo/__init__.py

@@ -0,0 +1,2 @@
+from .a2c_acktr import A2C_ACKTR
+from .ppo import PPO

a2c_ppo_acktr/algo/a2c_acktr.py

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+from a2c_ppo_acktr.algo.kfac import KFACOptimizer
+
+
+class A2C_ACKTR():
+    def __init__(self,
+                 actor_critic,
+                 value_loss_coef,
+                 entropy_coef,
+                 lr=None,
+                 eps=None,
+                 alpha=None,
+                 max_grad_norm=None,
+                 acktr=False):
+
+        self.actor_critic = actor_critic
+        self.acktr = acktr
+
+        self.value_loss_coef = value_loss_coef
+        self.entropy_coef = entropy_coef
+
+        self.max_grad_norm = max_grad_norm
+
+        if acktr:
+            self.optimizer = KFACOptimizer(actor_critic)
+        else:
+            self.optimizer = optim.RMSprop(
+                actor_critic.parameters(), lr, eps=eps, alpha=alpha)
+
+    def update(self, rollouts):
+        obs_shape = rollouts.obs.size()[2:]
+        action_shape = rollouts.actions.size()[-1]
+        num_steps, num_processes, _ = rollouts.rewards.size()
+
+        values, action_log_probs, dist_entropy, _ = self.actor_critic.evaluate_actions(
+            rollouts.obs[:-1].view(-1, *obs_shape),
+            rollouts.recurrent_hidden_states[0].view(
+                -1, self.actor_critic.recurrent_hidden_state_size),
+            rollouts.masks[:-1].view(-1, 1),
+            rollouts.actions.view(-1, action_shape))
+
+        values = values.view(num_steps, num_processes, 1)
+        action_log_probs = action_log_probs.view(num_steps, num_processes, 1)
+
+        advantages = rollouts.returns[:-1] - values
+        value_loss = advantages.pow(2).mean()
+
+        action_loss = -(advantages.detach() * action_log_probs).mean()
+
+        if self.acktr and self.optimizer.steps % self.optimizer.Ts == 0:
+            # Sampled fisher, see Martens 2014
+            self.actor_critic.zero_grad()
+            pg_fisher_loss = -action_log_probs.mean()
+
+            value_noise = torch.randn(values.size())
+            if values.is_cuda:
+                value_noise = value_noise.cuda()
+
+            sample_values = values + value_noise
+            vf_fisher_loss = -(values - sample_values.detach()).pow(2).mean()
+
+            fisher_loss = pg_fisher_loss + vf_fisher_loss
+            self.optimizer.acc_stats = True
+            fisher_loss.backward(retain_graph=True)
+            self.optimizer.acc_stats = False
+
+        self.optimizer.zero_grad()
+        (value_loss * self.value_loss_coef + action_loss -
+         dist_entropy * self.entropy_coef).backward()
+
+        if self.acktr == False:
+            nn.utils.clip_grad_norm_(self.actor_critic.parameters(),
+                                     self.max_grad_norm)
+
+        self.optimizer.step()
+
+        return value_loss.item(), action_loss.item(), dist_entropy.item()

a2c_ppo_acktr/algo/gail.py

@@ -0,0 +1,167 @@
+import h5py
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.data
+from torch import autograd
+
+from baselines.common.running_mean_std import RunningMeanStd
+
+
+class Discriminator(nn.Module):
+    def __init__(self, input_dim, hidden_dim, device):
+        super(Discriminator, self).__init__()
+
+        self.device = device
+
+        self.trunk = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim), nn.Tanh(),
+            nn.Linear(hidden_dim, hidden_dim), nn.Tanh(),
+            nn.Linear(hidden_dim, 1)).to(device)
+
+        self.trunk.train()
+
+        self.optimizer = torch.optim.Adam(self.trunk.parameters())
+
+        self.returns = None
+        self.ret_rms = RunningMeanStd(shape=())
+
+    def compute_grad_pen(self,
+                         expert_state,
+                         expert_action,
+                         policy_state,
+                         policy_action,
+                         lambda_=10):
+        alpha = torch.rand(expert_state.size(0), 1)
+        expert_data = torch.cat([expert_state, expert_action], dim=1)
+        policy_data = torch.cat([policy_state, policy_action], dim=1)
+
+        alpha = alpha.expand_as(expert_data).to(expert_data.device)
+
+        mixup_data = alpha * expert_data + (1 - alpha) * policy_data
+        mixup_data.requires_grad = True
+
+        disc = self.trunk(mixup_data)
+        ones = torch.ones(disc.size()).to(disc.device)
+        grad = autograd.grad(
+            outputs=disc,
+            inputs=mixup_data,
+            grad_outputs=ones,
+            create_graph=True,
+            retain_graph=True,
+            only_inputs=True)[0]
+
+        grad_pen = lambda_ * (grad.norm(2, dim=1) - 1).pow(2).mean()
+        return grad_pen
+
+    def update(self, expert_loader, rollouts, obsfilt=None):
+        self.train()
+
+        policy_data_generator = rollouts.feed_forward_generator(
+            None, mini_batch_size=expert_loader.batch_size)
+
+        loss = 0
+        n = 0
+        for expert_batch, policy_batch in zip(expert_loader,
+                                              policy_data_generator):
+            policy_state, policy_action = policy_batch[0], policy_batch[2]
+            policy_d = self.trunk(
+                torch.cat([policy_state, policy_action], dim=1))
+
+            expert_state, expert_action = expert_batch
+            expert_state = obsfilt(expert_state.numpy(), update=False)
+            expert_state = torch.FloatTensor(expert_state).to(self.device)
+            expert_action = expert_action.to(self.device)
+            expert_d = self.trunk(
+                torch.cat([expert_state, expert_action], dim=1))
+
+            expert_loss = F.binary_cross_entropy_with_logits(
+                expert_d,
+                torch.ones(expert_d.size()).to(self.device))
+            policy_loss = F.binary_cross_entropy_with_logits(
+                policy_d,
+                torch.zeros(policy_d.size()).to(self.device))
+
+            gail_loss = expert_loss + policy_loss
+            grad_pen = self.compute_grad_pen(expert_state, expert_action,
+                                             policy_state, policy_action)
+
+            loss += (gail_loss + grad_pen).item()
+            n += 1
+
+            self.optimizer.zero_grad()
+            (gail_loss + grad_pen).backward()
+            self.optimizer.step()
+        return loss / n
+
+    def predict_reward(self, state, action, gamma, masks, update_rms=True):
+        with torch.no_grad():
+            self.eval()
+            d = self.trunk(torch.cat([state, action], dim=1))
+            s = torch.sigmoid(d)
+            reward = s.log() - (1 - s).log()
+            if self.returns is None:
+                self.returns = reward.clone()
+
+            if update_rms:
+                self.returns = self.returns * masks * gamma + reward
+                self.ret_rms.update(self.returns.cpu().numpy())
+
+            return reward / np.sqrt(self.ret_rms.var[0] + 1e-8)
+
+
+class ExpertDataset(torch.utils.data.Dataset):
+    def __init__(self, file_name, num_trajectories=4, subsample_frequency=20):
+        all_trajectories = torch.load(file_name)
+
+        perm = torch.randperm(all_trajectories['states'].size(0))
+        idx = perm[:num_trajectories]
+
+        self.trajectories = {}
+
+        # See https://github.com/pytorch/pytorch/issues/14886
+        # .long() for fixing bug in torch v0.4.1
+        start_idx = torch.randint(
+            0, subsample_frequency, size=(num_trajectories, )).long()
+
+        for k, v in all_trajectories.items():
+            data = v[idx]
+
+            if k != 'lengths':
+                samples = []
+                for i in range(num_trajectories):
+                    samples.append(data[i, start_idx[i]::subsample_frequency])
+                self.trajectories[k] = torch.stack(samples)
+            else:
+                self.trajectories[k] = data // subsample_frequency
+
+        self.i2traj_idx = {}
+        self.i2i = {}
+
+        self.length = self.trajectories['lengths'].sum().item()
+
+        traj_idx = 0
+        i = 0
+
+        self.get_idx = []
+
+        for j in range(self.length):
+
+            while self.trajectories['lengths'][traj_idx].item() <= i:
+                i -= self.trajectories['lengths'][traj_idx].item()
+                traj_idx += 1
+
+            self.get_idx.append((traj_idx, i))
+
+            i += 1
+
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, i):
+        traj_idx, i = self.get_idx[i]
+
+        return self.trajectories['states'][traj_idx][i], self.trajectories[
+            'actions'][traj_idx][i]