iicd/xautodl
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update LFNA version 1.0

Browse Source
This commit is contained in:
D-X-Y
2021-05-13 21:33:34 +08:00
parent 3d3a04705f
commit cfabd05de8
11 changed files with 340 additions and 299 deletions
Download Patch File Download Diff File Expand all files Collapse all files

120
lib/xlayers/super_rl_actor.py
View File

@@ -1,120 +0,0 @@
#####################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
#####################################################
# DISABLED / NOT-FINISHED
#####################################################
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Callable
import spaces
from .super_container import SuperSequential
from .super_linear import SuperLinear
class SuperActor(SuperModule):
"""A Actor in RL."""
def _distribution(self, obs):
raise NotImplementedError
def _log_prob_from_distribution(self, pi, act):
raise NotImplementedError
def forward_candidate(self, **kwargs):
return self.forward_raw(**kwargs)
def forward_raw(self, obs, act=None):
# Produce action distributions for given observations, and
# optionally compute the log likelihood of given actions under
# those distributions.
pi = self._distribution(obs)
logp_a = None
if act is not None:
logp_a = self._log_prob_from_distribution(pi, act)
return pi, logp_a
class SuperLfnaMetaMLP(SuperModule):
def __init__(self, obs_dim, hidden_sizes, act_cls):
super(SuperLfnaMetaMLP).__init__()
self.delta_net = SuperSequential(
SuperLinear(obs_dim, hidden_sizes[0]),
act_cls(),
SuperLinear(hidden_sizes[0], hidden_sizes[1]),
act_cls(),
SuperLinear(hidden_sizes[1], 1),
)
class SuperLfnaMetaMLP(SuperModule):
def __init__(self, obs_dim, act_dim, hidden_sizes, act_cls):
super(SuperLfnaMetaMLP).__init__()
log_std = -0.5 * np.ones(act_dim, dtype=np.float32)
self.log_std = torch.nn.Parameter(torch.as_tensor(log_std))
self.mu_net = SuperSequential(
SuperLinear(obs_dim, hidden_sizes[0]),
act_cls(),
SuperLinear(hidden_sizes[0], hidden_sizes[1]),
act_cls(),
SuperLinear(hidden_sizes[1], act_dim),
)
def _distribution(self, obs):
mu = self.mu_net(obs)
std = torch.exp(self.log_std)
return Normal(mu, std)
def _log_prob_from_distribution(self, pi, act):
return pi.log_prob(act).sum(axis=-1)
def forward_candidate(self, **kwargs):
return self.forward_raw(**kwargs)
def forward_raw(self, obs, act=None):
# Produce action distributions for given observations, and
# optionally compute the log likelihood of given actions under
# those distributions.
pi = self._distribution(obs)
logp_a = None
if act is not None:
logp_a = self._log_prob_from_distribution(pi, act)
return pi, logp_a
class SuperMLPGaussianActor(SuperModule):
def __init__(self, obs_dim, act_dim, hidden_sizes, act_cls):
super(SuperMLPGaussianActor).__init__()
log_std = -0.5 * np.ones(act_dim, dtype=np.float32)
self.log_std = torch.nn.Parameter(torch.as_tensor(log_std))
self.mu_net = SuperSequential(
SuperLinear(obs_dim, hidden_sizes[0]),
act_cls(),
SuperLinear(hidden_sizes[0], hidden_sizes[1]),
act_cls(),
SuperLinear(hidden_sizes[1], act_dim),
)
def _distribution(self, obs):
mu = self.mu_net(obs)
std = torch.exp(self.log_std)
return Normal(mu, std)
def _log_prob_from_distribution(self, pi, act):
return pi.log_prob(act).sum(axis=-1)
def forward_candidate(self, **kwargs):
return self.forward_raw(**kwargs)
def forward_raw(self, obs, act=None):
# Produce action distributions for given observations, and
# optionally compute the log likelihood of given actions under
# those distributions.
pi = self._distribution(obs)
logp_a = None
if act is not None:
logp_a = self._log_prob_from_distribution(pi, act)
return pi, logp_a

9
lib/xlayers/super_transformer.py
View File

@@ -42,6 +42,7 @@ class SuperTransformerEncoderLayer(SuperModule):
qkv_bias: BoolSpaceType = False,
mlp_hidden_multiplier: IntSpaceType = 4,
drop: Optional[float] = None,
norm_affine: bool = True,
act_layer: Callable[[], nn.Module] = nn.GELU,
order: LayerOrder = LayerOrder.PreNorm,
):
@@ -62,19 +63,19 @@ class SuperTransformerEncoderLayer(SuperModule):
drop=drop,
)
if order is LayerOrder.PreNorm:
self.norm1 = SuperLayerNorm1D(d_model)
self.norm1 = SuperLayerNorm1D(d_model, elementwise_affine=norm_affine)
self.mha = mha
self.drop1 = nn.Dropout(drop or 0.0)
self.norm2 = SuperLayerNorm1D(d_model)
self.norm2 = SuperLayerNorm1D(d_model, elementwise_affine=norm_affine)
self.mlp = mlp
self.drop2 = nn.Dropout(drop or 0.0)
elif order is LayerOrder.PostNorm:
self.mha = mha
self.drop1 = nn.Dropout(drop or 0.0)
self.norm1 = SuperLayerNorm1D(d_model)
self.norm1 = SuperLayerNorm1D(d_model, elementwise_affine=norm_affine)
self.mlp = mlp
self.drop2 = nn.Dropout(drop or 0.0)
self.norm2 = SuperLayerNorm1D(d_model)
self.norm2 = SuperLayerNorm1D(d_model, elementwise_affine=norm_affine)
else:
raise ValueError("Unknown order: {:}".format(order))
self._order = order

5
lib/xlayers/weight_init.py
View File

@@ -60,4 +60,7 @@ def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
>>> w = torch.empty(3, 5)
>>> nn.init.trunc_normal_(w)
"""
return _no_grad_trunc_normal_(tensor, mean, std, a, b)
if isinstance(tensor, list):
return [_no_grad_trunc_normal_(x, mean, std, a, b) for x in tensor]
else:
return _no_grad_trunc_normal_(tensor, mean, std, a, b)