update

Author: Bennyoooo

.gitignore

@@ -1,2 +1,3 @@
 **/data
-*.pyc
+*.pyc
+metaworld

cs285pkg/cs285/agents/__pycache__/__init__.cpython-36.pyc

@@ -4,7 +4,7 @@
 from cs285.policies.MLP_policy import MLPPolicyPG
 from cs285.infrastructure.replay_buffer import ReplayBuffer
 from cs285.infrastructure.utils import normalize
-from cs285.policies.PSPPolicy import PSPPolicy
+
 
 class PGAgent(BaseAgent):
     def __init__(self, env, agent_params):
@@ -29,17 +29,6 @@ def __init__(self, env, agent_params):
             nn_baseline=self.agent_params["nn_baseline"],
         )
 
-        # self.actor = PSPPolicy(
-        #     self.agent_params["ac_dim"],
-        #     self.agent_params["ob_dim"],
-        #     self.agent_params["n_layers"],
-        #     self.agent_params["size"],
-        #     period=agent_params['period'],
-        #     discrete=self.agent_params["discrete"],
-        #     learning_rate=self.agent_params["learning_rate"],
-        #     nn_baseline=self.agent_params["nn_baseline"],
-        # )
-
         # replay buffer
         self.replay_buffer = ReplayBuffer(1000000)
 

cs285pkg/cs285/agents/__pycache__/pg_agent.cpython-36.pyc

@@ -9,22 +9,56 @@
 class BinaryHashLinear(nn.Module):
     def __init__(self, n_in, n_out, period, key_pick='hash', learn_key=True):
         super(BinaryHashLinear, self).__init__()
-        # self.key_pick = key_pick
-        # w = nn.init.xavier_normal_(torch.empty(n_in, n_out))
-        # rand_01 = np.random.binomial(p=.5, n=1, size=(n_in, period)).astype(np.float32)
-        # o = torch.from_numpy(rand_01*2 - 1)
-        # self.n_in = n_in
-        # self.n_out = n_out
-        #
-        # self.w = nn.Parameter(w)
-        # self.bias = nn.Parameter(torch.zeros(n_out))
-        # self.o = nn.Parameter(o)
-        self.linear = nn.Linear(n_in, n_out)
-        # if not learn_key:
-        #     self.o.requires_grad = False
-
-    def forward(self, x, time=0):
-        # o = self.o[:, int(time)]
-        # m = x*o
-        r = self.linear(x)
+        self.key_pick = key_pick
+        w = nn.init.xavier_normal_(torch.empty(n_in, n_out))
+        rand_01 = np.random.binomial(p=.5, n=1, size=(n_in, period)).astype(np.float32)
+        o = torch.from_numpy(rand_01*2 - 1)
+
+        self.w = nn.Parameter(w)
+        self.bias = nn.Parameter(torch.zeros(n_out))
+        self.o = nn.Parameter(o)
+        if not learn_key:
+            self.o.requires_grad = False
+
+    def forward(self, x, time):
+        o = self.o[:, int(time)]
+        m = x*o
+        r = torch.mm(m, self.w)
         return r
+
+
+class HashLinear(nn.Module):
+    '''Complex layer with complex diagonal contexts'''
+    def __init__(self, n_in, n_out, period=2, key_pick='hash', learn_key=True):
+        super(HashLinear, self).__init__()
+        self.key_pick = key_pick
+        w_r = nn.init.xavier_normal_(torch.empty(n_in, n_out))
+        w_phi = torch.Tensor(n_in, n_out).uniform_(-np.pi, np.pi)
+        o_r = torch.ones(period, n_in)
+        o_phi = torch.Tensor(period, n_in).uniform_(-np.pi, np.pi)
+
+        self.w = nn.Parameter(torch.stack(from_polar(w_r, w_phi)))
+        self.bias = nn.Parameter(torch.zeros(n_out))
+        self.o = nn.Parameter(torch.stack(from_polar(o_r, o_phi)))
+        if not learn_key:
+            self.o.requires_grad = False
+
+    def forward(self, x_a, x_b, time):
+        net_time = int(time) % self.o.shape[1]
+        o = self.o[:, net_time]
+        o_a = o[0].unsqueeze(0)
+        o_b = o[1].unsqueeze(0)
+        m_a = x_a*o_a - x_b*o_b
+        m_b = x_b*o_a + x_a*o_b
+
+        w_a = self.w[0]
+        w_b = self.w[1]
+        r_a = torch.mm(m_a, w_a) - torch.mm(m_b, w_b)
+        r_b = torch.mm(m_b, w_a) + torch.mm(m_a, w_b)
+        return r_a + self.bias, r_b
+
+
+def from_polar(r, phi):
+    a = r*torch.cos(phi)
+    b = r*torch.sin(phi)
+    return a, b

cs285pkg/cs285/agents/pg_agent.py

@@ -0,0 +1,27 @@
+import numpy as np
+import scipy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.modules.utils import _pair
+
+
+class BinaryHashLinear(nn.Module):
+    def __init__(self, n_in, n_out, period, key_pick='hash', learn_key=True):
+        super(BinaryHashLinear, self).__init__()
+        self.key_pick = key_pick
+        w = nn.init.xavier_normal_(torch.empty(n_in, n_out))
+        rand_01 = np.random.binomial(p=.5, n=1, size=(n_in, period)).astype(np.float32)
+        o = torch.from_numpy(rand_01*2 - 1)
+
+        self.w = nn.Parameter(w)
+        self.bias = nn.Parameter(torch.zeros(n_out))
+        self.o = nn.Parameter(o)
+        if not learn_key:
+            self.o.requires_grad = False
+
+    def forward(self, x, time):
+        o = self.o[:, int(time)]
+        m = x*o
+        r = torch.mm(m, self.w)
+        return r

cs285pkg/cs285/infrastructure/__pycache__/__init__.cpython-36.pyc

@@ -1,5 +1,5 @@
-#from cs285.infrastructure.psp_layer import *
-from cs285.infrastructure.psp_layer2 import *
+from cs285.infrastructure.psp_layer import *
+
 
 class HashNet(nn.Module):
     def __init__(self, input_dim, output_dim, layer_size,
@@ -29,12 +29,10 @@ def forward(self, x, time):
                 r = self.activation(r)
             r = layer(r, time)
             preactivations.append(r)
-        r = nn.Identity()(r)
 
         return r, None, preactivations
 
 
-
 class ComplexHashNet(HashNet):
     def forward(self, x, time):
         preactivations = []
@@ -43,26 +41,7 @@ def forward(self, x, time):
             if layer_i > 0:
                 r_a = self.activation(r_a)
                 r_b = self.activation(r_b)
-            r_a = layer(r_a, time)
-            r_b = layer(r_b, time)
+            r_a, r_b = layer(r_a, r_b, time)
             preactivations.append(r_a)
-            preactivations.append(r_b)
-        return r_a, r_b, preactivations
-
 
-class MLP(nn.Module):
-
-    def __init__(self, input_size, output_size, n_layers, size, activation, output_activation):
-        super(MLP, self).__init__()
-        self.n_layers = n_layers
-        self.linears = nn.ModuleList([nn.Linear(input_size, size)])
-        self.linears.extend([nn.Linear(size, size) for i in range(0, self.n_layers - 1)])
-        self.linears.append(nn.Linear(size, output_size))
-        self.activation = activation
-        self.output_activation = output_activation
-
-    def forward(self, x):
-        for i in range(self.n_layers):
-            x = self.activation(self.linears[i](x))
-        mean = self.output_activation(self.linears[self.n_layers](x))
-        return mean
+        return r_a, r_b, preactivations

cs285pkg/cs285/infrastructure/psp_layer.py

@@ -2,7 +2,6 @@
 
 import torch
 from torch import nn
-from cs285.infrastructure.psp_layer import BinaryHashLinear
 
 Activation = Union[str, nn.Module]
 

cs285pkg/cs285/infrastructure/psp_layer.py.bak

@@ -129,6 +129,7 @@ def run_training_loop(
 
         for itr in range(n_iter):
             print("\n\n********** Iteration %i ************" % itr)
+            self.agent.update_time(0)
 
             # decide if videos should be rendered/logged at this iteration
             if (
@@ -239,11 +240,12 @@ def run_second_task_loop(
             # log/save
             if self.logvideo or self.logmetrics:
                 # perform logging
+                self.agent.update_time(0)
                 print("\nBeginning logging procedure...")
                 self.perform_logging(
                     itr, paths, eval_policy, train_video_paths, train_logs
                 )
-
+                self.agent.update_time(1)
                 if self.params["save_params"]:
                     self.agent.save(
                         "{}/agent_itr_{}.pt".format(self.params["logdir"], itr)
@@ -328,30 +330,10 @@ def perform_logging(self, itr, paths, eval_policy, train_video_paths, all_logs):
         eval_paths, eval_envsteps_this_batch = utils.sample_trajectories(
             self.env, eval_policy, self.params["eval_batch_size"], self.params["ep_len"]
         )
-
-        # save eval rollouts as videos in tensorboard event file
-        if self.logvideo and train_video_paths != None:
-            print("\nCollecting video rollouts eval")
-            eval_video_paths = utils.sample_n_trajectories(
-                self.env, eval_policy, MAX_NVIDEO, MAX_VIDEO_LEN, True
-            )
-
-            # save train/eval videos
-            print("\nSaving train rollouts as videos...")
-            self.logger.log_paths_as_videos(
-                train_video_paths,
-                itr,
-                fps=self.fps,
-                max_videos_to_save=MAX_NVIDEO,
-                video_title="train_rollouts",
-            )
-            self.logger.log_paths_as_videos(
-                eval_video_paths,
-                itr,
-                fps=self.fps,
-                max_videos_to_save=MAX_NVIDEO,
-                video_title="eval_rollouts",
-            )
+        self.agent.update_time(1)
+        eval_paths2, eval_envsteps_this_batch2 = utils.sample_trajectories(
+            self.env2, eval_policy, self.params["eval_batch_size"], self.params["ep_len"]
+        )
 
         #######################
 
@@ -360,18 +342,22 @@ def perform_logging(self, itr, paths, eval_policy, train_video_paths, all_logs):
             # returns, for logging
             train_returns = [path["reward"].sum() for path in paths]
             eval_returns = [eval_path["reward"].sum() for eval_path in eval_paths]
+            eval_returns2 = [eval_path["reward"].sum() for eval_path in eval_paths2]
 
             # episode lengths, for logging
             train_ep_lens = [len(path["reward"]) for path in paths]
             eval_ep_lens = [len(eval_path["reward"]) for eval_path in eval_paths]
+            eval_ep_lens2 = [len(eval_path["reward"]) for eval_path in eval_paths2]
 
             # decide what to log
             logs = OrderedDict()
             logs["Eval_AverageReturn"] = np.mean(eval_returns)
+            logs["Eval_AverageReturn2"] = np.mean(eval_returns2)
             logs["Eval_StdReturn"] = np.std(eval_returns)
             logs["Eval_MaxReturn"] = np.max(eval_returns)
             logs["Eval_MinReturn"] = np.min(eval_returns)
             logs["Eval_AverageEpLen"] = np.mean(eval_ep_lens)
+            logs["Eval_AverageEpLen2"] = np.mean(eval_ep_lens2)
 
             logs["Train_AverageReturn"] = np.mean(train_returns)
             logs["Train_StdReturn"] = np.std(train_returns)

cs285pkg/cs285/infrastructure/psp_layer2.py

@@ -11,9 +11,6 @@
 from cs285.infrastructure import pytorch_util as ptu
 from cs285.policies.base_policy import BasePolicy
 from cs285.infrastructure.utils import normalize
-from cs285.infrastructure.psp_net import RealHashNet
-from cs285.infrastructure.psp_layer import BinaryHashLinear
-
 
 
 class MLPPolicy(BasePolicy, nn.Module, metaclass=abc.ABCMeta):
@@ -60,7 +57,6 @@ def __init__(
                 n_layers=self.n_layers,
                 size=self.size,
             )
-            # self.mean_net = RealHashNet(self.ob_dim, self.ac_dim, self.size, torch.tanh, self.n_layers, 2, 'hash', BinaryHashLinear)
             self.logstd = nn.Parameter(
                 torch.zeros(self.ac_dim, dtype=torch.float32, device=ptu.device)
             )

cs285pkg/cs285/infrastructure/psp_net.py

@@ -11,8 +11,7 @@
 from cs285.infrastructure import pytorch_util as ptu
 from cs285.policies.base_policy import BasePolicy
 from cs285.infrastructure.utils import normalize
-from cs285.infrastructure.psp_net import RealHashNet, ComplexHashNet, MLP
-from cs285.infrastructure.psp_layer import *
+from cs285.infrastructure.psp_net import RealHashNet, BinaryHashLinear, ComplexHashNet, HashLinear
 
 
 class PSPPolicy(BasePolicy, nn.Module, metaclass=abc.ABCMeta):
@@ -40,13 +39,9 @@ def __init__(
         self.learning_rate = learning_rate
         self.training = training
         self.period = period
-        self.mean_net = MLP(self.ob_dim, self.ac_dim, n_layers, size, nn.Tanh(), nn.Identity())
-        self.mean_net2 = ptu.build_mlp(
-                input_size=self.ob_dim,
-                output_size=self.ac_dim,
-                n_layers=self.n_layers,
-                size=self.size,
-            )
+        # self.mean_net = RealHashNet(self.ob_dim, self.ac_dim, self.size, torch.tanh, self.n_layers, self.period, 'hash', BinaryHashLinear)
+        self.mean_net = ComplexHashNet(self.ob_dim, self.ac_dim, self.size, torch.tanh, self.n_layers, self.period, 'hash',
+                                    HashLinear)
         self.logstd = nn.Parameter(
             torch.zeros(self.ac_dim, dtype=torch.float32, device=ptu.device)
         )
@@ -56,7 +51,6 @@ def __init__(
             itertools.chain([self.logstd], self.mean_net.parameters()),
             self.learning_rate,
         )
-        self.a = self.mean_net.parameters()
 
     def update_time(self, time):
         self.time = time
@@ -86,10 +80,8 @@ def update(self, observations, acs_na, adv_n=None, acs_labels_na=None,
                qvals=None):
         observations = ptu.from_numpy(observations)
         actions = ptu.from_numpy(acs_na)
-        adv_n = ptu.from_numpy(adv_n)
-
         action_distribution = self(observations)
-        loss = - action_distribution.log_prob(actions) * adv_n
+        loss = -action_distribution.log_prob(actions) * ptu.from_numpy(adv_n)
         loss = loss.mean()
 
         self.optimizer.zero_grad()
@@ -106,12 +98,11 @@ def update(self, observations, acs_na, adv_n=None, acs_labels_na=None,
     # return more flexible objects, such as a
     # `torch.distributions.Distribution` object. It's up to you!
     def forward(self, observation: torch.FloatTensor):
-        batch_mean = self.mean_net(observation)
-        # batch_mean2 = self.mean_net2(observation)
+        batch_mean = self.mean_net(observation, self.time)[0]
         scale_tril = torch.diag(torch.exp(self.logstd))
         batch_dim = batch_mean.shape[0]
         batch_scale_tril = scale_tril.repeat(batch_dim, 1, 1)
         action_distribution = distributions.MultivariateNormal(
             batch_mean, scale_tril=batch_scale_tril,
         )
-        return action_distribution
+        return action_distribution