From d3bc20b727b1e4c58644a5be303b8f9be0510549 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Mon, 12 Oct 2020 15:00:42 +0200 Subject: [PATCH] [RLlib] ConvTranspose2D module (#11231) --- rllib/BUILD | 21 +- rllib/agents/dreamer/dreamer_model.py | 818 +++++++++--------- rllib/agents/dreamer/utils.py | 17 +- rllib/models/tests/test_attention_nets.py | 2 +- .../tests/test_convtranspose2d_stack.py | 64 ++ rllib/models/tf/layers/noisy_layer.py | 12 +- rllib/models/torch/misc.py | 13 + .../torch/modules/convtranspose2d_stack.py | 77 ++ rllib/models/torch/modules/noisy_layer.py | 12 +- rllib/models/utils.py | 35 + rllib/tests/data/images/obstacle_tower.png | Bin 0 -> 8737 bytes rllib/utils/framework.py | 3 +- 12 files changed, 621 insertions(+), 453 deletions(-) create mode 100644 rllib/models/tests/test_convtranspose2d_stack.py create mode 100644 rllib/models/torch/modules/convtranspose2d_stack.py create mode 100644 rllib/tests/data/images/obstacle_tower.png diff --git a/rllib/BUILD b/rllib/BUILD index 7e48f70e2..a1afef2f6 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -1040,13 +1040,6 @@ py_test( # Tag: models # -------------------------------------------------------------------- -py_test( - name = "test_distributions", - tags = ["models"], - size = "medium", - srcs = ["models/tests/test_distributions.py"] -) - py_test( name = "test_attention_nets", tags = ["models"], @@ -1054,6 +1047,20 @@ py_test( srcs = ["models/tests/test_attention_nets.py"] ) +py_test( + name = "test_convtranspose2d_stack", + tags = ["models"], + size = "small", + data = glob(["tests/data/images/obstacle_tower.png"]), + srcs = ["models/tests/test_convtranspose2d_stack.py"] +) + +py_test( + name = "test_distributions", + tags = ["models"], + size = "medium", + srcs = ["models/tests/test_distributions.py"] +) # -------------------------------------------------------------------- # Evaluation components diff --git a/rllib/agents/dreamer/dreamer_model.py b/rllib/agents/dreamer/dreamer_model.py index 2daeee643..5483f664f 100644 --- a/rllib/agents/dreamer/dreamer_model.py +++ b/rllib/agents/dreamer/dreamer_model.py @@ -1,5 +1,6 @@ import numpy as np from typing import Any, List, Tuple +from ray.rllib.models.torch.modules.reshape import Reshape from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.utils.framework import try_import_torch from ray.rllib.utils.framework import TensorType @@ -12,203 +13,187 @@ if torch: ActFunc = Any + # Encoder, part of PlaNET -if torch: +class ConvEncoder(nn.Module): + """Standard Convolutional Encoder for Dreamer. This encoder is used + to encode images frm an enviornment into a latent state for the + RSSM model in PlaNET. + """ - class ConvEncoder(nn.Module): - """Standard Convolutional Encoder for Dreamer. This encoder is used - to encode images frm an enviornment into a latent state for the - RSSM model in PlaNET. - """ + def __init__(self, + depth: int = 32, + act: ActFunc = None, + shape: Tuple[int] = (3, 64, 64)): + """Initializes Conv Encoder - def __init__(self, - depth: int = 32, - act: ActFunc = None, - shape: List = [3, 64, 64]): - """Initializes Conv Encoder - - Args: + Args: depth (int): Number of channels in the first conv layer act (Any): Activation for Encoder, default ReLU shape (List): Shape of observation input - """ - super().__init__() - self.act = act - if not act: - self.act = nn.ReLU - self.depth = depth - self.shape = shape + """ + super().__init__() + self.act = act + if not act: + self.act = nn.ReLU + self.depth = depth + self.shape = shape - init_channels = self.shape[0] - self.layers = [ - Conv2d(init_channels, self.depth, 4, stride=2), - self.act(), - Conv2d(self.depth, 2 * self.depth, 4, stride=2), - self.act(), - Conv2d(2 * self.depth, 4 * self.depth, 4, stride=2), - self.act(), - Conv2d(4 * self.depth, 8 * self.depth, 4, stride=2), - self.act(), - ] - self.model = nn.Sequential(*self.layers) + init_channels = self.shape[0] + self.layers = [ + Conv2d(init_channels, self.depth, 4, stride=2), + self.act(), + Conv2d(self.depth, 2 * self.depth, 4, stride=2), + self.act(), + Conv2d(2 * self.depth, 4 * self.depth, 4, stride=2), + self.act(), + Conv2d(4 * self.depth, 8 * self.depth, 4, stride=2), + self.act(), + ] + self.model = nn.Sequential(*self.layers) - def forward(self, x): - # Flatten to [batch*horizon, 3, 64, 64] in loss function - orig_shape = list(x.size()) - x = x.view(-1, *(orig_shape[-3:])) - x = self.model(x) + def forward(self, x): + # Flatten to [batch*horizon, 3, 64, 64] in loss function + orig_shape = list(x.size()) + x = x.view(-1, *(orig_shape[-3:])) + x = self.model(x) - new_shape = orig_shape[:-3] + [32 * self.depth] - x = x.view(*new_shape) - return x - - -if torch: - - class Reshape(nn.Module): - """Standard module that reshapes/views a tensor - """ - - def __init__(self, shape: List): - super().__init__() - self.shape = shape - - def forward(self, x): - return x.view(*self.shape) + new_shape = orig_shape[:-3] + [32 * self.depth] + x = x.view(*new_shape) + return x # Decoder, part of PlaNET -if torch: +class ConvDecoder(nn.Module): + """Standard Convolutional Decoder for Dreamer. - class ConvDecoder(nn.Module): - """Standard Convolutional Decoder for Dreamer. This decoder is used - to decoder images from the latent state generated by the transition - dynamics model. This is used in calulating loss and logging gifs for - imagine trajectories. - """ + This decoder is used to decode images from the latent state generated + by the transition dynamics model. This is used in calculating loss and + logging gifs for imagined trajectories. + """ - def __init__(self, - input_size: int, - depth: int = 32, - act: ActFunc = None, - shape: List = [3, 64, 64]): - """Initializes Conv Decoder + def __init__(self, + input_size: int, + depth: int = 32, + act: ActFunc = None, + shape: Tuple[int] = (3, 64, 64)): + """Initializes a ConvDecoder instance. - Args: - input_size (int): Input size, usually feature size output from RSSM + Args: + input_size (int): Input size, usually feature size output from + RSSM. depth (int): Number of channels in the first conv layer act (Any): Activation for Encoder, default ReLU shape (List): Shape of observation input - """ - super().__init__() - self.act = act - if not act: - self.act = nn.ReLU - self.depth = depth - self.shape = shape + """ + super().__init__() + self.act = act + if not act: + self.act = nn.ReLU + self.depth = depth + self.shape = shape - self.layers = [ - Linear(input_size, 32 * self.depth), - Reshape((-1, 32 * self.depth, 1, 1)), - ConvTranspose2d(32 * self.depth, 4 * self.depth, 5, stride=2), - self.act(), - ConvTranspose2d(4 * self.depth, 2 * self.depth, 5, stride=2), - self.act(), - ConvTranspose2d(2 * self.depth, self.depth, 6, stride=2), - self.act(), - ConvTranspose2d(self.depth, self.shape[0], 6, stride=2), - ] - self.model = nn.Sequential(*self.layers) + self.layers = [ + Linear(input_size, 32 * self.depth), + Reshape([-1, 32 * self.depth, 1, 1]), + ConvTranspose2d(32 * self.depth, 4 * self.depth, 5, stride=2), + self.act(), + ConvTranspose2d(4 * self.depth, 2 * self.depth, 5, stride=2), + self.act(), + ConvTranspose2d(2 * self.depth, self.depth, 6, stride=2), + self.act(), + ConvTranspose2d(self.depth, self.shape[0], 6, stride=2), + ] + self.model = nn.Sequential(*self.layers) - def forward(self, x): - # x is [batch, hor_length, input_size] - orig_shape = list(x.size()) - x = self.model(x) + def forward(self, x): + # x is [batch, hor_length, input_size] + orig_shape = list(x.size()) + x = self.model(x) - reshape_size = orig_shape[:-1] + self.shape - mean = x.view(*reshape_size) + reshape_size = orig_shape[:-1] + self.shape + mean = x.view(*reshape_size) - # Equivalent to making a multivariate diag - return td.Independent(td.Normal(mean, 1), len(self.shape)) + # Equivalent to making a multivariate diag + return td.Independent(td.Normal(mean, 1), len(self.shape)) # Reward Model (PlaNET), and Value Function -if torch: +class DenseDecoder(nn.Module): + """FC network that outputs a distribution for calculating log_prob. - class DenseDecoder(nn.Module): - """Fully Connected network that outputs a distribution for calculating log_prob - later in DreamerLoss - """ + Used later in DreamerLoss. + """ - def __init__(self, - input_size: int, - output_size: int, - layers: int, - units: int, - dist: str = "normal", - act: ActFunc = None): - """Initializes FC network + def __init__(self, + input_size: int, + output_size: int, + layers: int, + units: int, + dist: str = "normal", + act: ActFunc = None): + """Initializes FC network - Args: + Args: input_size (int): Input size to network output_size (int): Output size to network layers (int): Number of layers in network units (int): Size of the hidden layers - dist (str): Output distribution, parameterized by FC output logits + dist (str): Output distribution, parameterized by FC output + logits. act (Any): Activation function - """ - super().__init__() - self.layrs = layers - self.units = units - self.act = act - if not act: - self.act = nn.ELU - self.dist = dist - self.input_size = input_size - self.output_size = output_size - self.layers = [] - cur_size = input_size - for _ in range(self.layrs): - self.layers.extend([Linear(cur_size, self.units), self.act()]) - cur_size = units - self.layers.append(Linear(cur_size, output_size)) - self.model = nn.Sequential(*self.layers) + """ + super().__init__() + self.layrs = layers + self.units = units + self.act = act + if not act: + self.act = nn.ELU + self.dist = dist + self.input_size = input_size + self.output_size = output_size + self.layers = [] + cur_size = input_size + for _ in range(self.layrs): + self.layers.extend([Linear(cur_size, self.units), self.act()]) + cur_size = units + self.layers.append(Linear(cur_size, output_size)) + self.model = nn.Sequential(*self.layers) - def forward(self, x): - x = self.model(x) - if self.output_size == 1: - x = torch.squeeze(x) - if self.dist == "normal": - output_dist = td.Normal(x, 1) - elif self.dist == "binary": - output_dist = td.Bernoulli(logits=x) - else: - raise NotImplementedError("Distribution type not implemented!") - return td.Independent(output_dist, 0) + def forward(self, x): + x = self.model(x) + if self.output_size == 1: + x = torch.squeeze(x) + if self.dist == "normal": + output_dist = td.Normal(x, 1) + elif self.dist == "binary": + output_dist = td.Bernoulli(logits=x) + else: + raise NotImplementedError("Distribution type not implemented!") + return td.Independent(output_dist, 0) # Represents dreamer policy -if torch: +class ActionDecoder(nn.Module): + """ActionDecoder is the policy module in Dreamer. - class ActionDecoder(nn.Module): - """ActionDecoder is the policy module in Dreamer. It outputs a distribution - parameterized by mean and std, later to be transformed by a custom - TanhBijector in utils.py for Dreamer. - """ + It outputs a distribution parameterized by mean and std, later to be + transformed by a custom TanhBijector in utils.py for Dreamer. + """ - def __init__(self, - input_size: int, - action_size: int, - layers: int, - units: int, - dist: str = "tanh_normal", - act: ActFunc = None, - min_std: float = 1e-4, - init_std: float = 5.0, - mean_scale: float = 5.0): - """Initializes Policy + def __init__(self, + input_size: int, + action_size: int, + layers: int, + units: int, + dist: str = "tanh_normal", + act: ActFunc = None, + min_std: float = 1e-4, + init_std: float = 5.0, + mean_scale: float = 5.0): + """Initializes Policy - Args: + Args: input_size (int): Input size to network action_size (int): Action space size layers (int): Number of layers in network @@ -218,342 +203,335 @@ if torch: min_std (float): Minimum std for output distribution init_std (float): Intitial std mean_scale (float): Augmenting mean output from FC network - """ - super().__init__() - self.layrs = layers - self.units = units - self.dist = dist - self.act = act - if not act: - self.act = nn.ReLU - self.min_std = min_std - self.init_std = init_std - self.mean_scale = mean_scale - self.action_size = action_size + """ + super().__init__() + self.layrs = layers + self.units = units + self.dist = dist + self.act = act + if not act: + self.act = nn.ReLU + self.min_std = min_std + self.init_std = init_std + self.mean_scale = mean_scale + self.action_size = action_size - self.layers = [] - self.softplus = nn.Softplus() + self.layers = [] + self.softplus = nn.Softplus() - # MLP Construction - cur_size = input_size - for _ in range(self.layrs): - self.layers.extend([Linear(cur_size, self.units), self.act()]) - cur_size = self.units - if self.dist == "tanh_normal": - self.layers.append(Linear(cur_size, 2 * action_size)) - elif self.dist == "onehot": - self.layers.append(Linear(cur_size, action_size)) - self.model = nn.Sequential(*self.layers) + # MLP Construction + cur_size = input_size + for _ in range(self.layrs): + self.layers.extend([Linear(cur_size, self.units), self.act()]) + cur_size = self.units + if self.dist == "tanh_normal": + self.layers.append(Linear(cur_size, 2 * action_size)) + elif self.dist == "onehot": + self.layers.append(Linear(cur_size, action_size)) + self.model = nn.Sequential(*self.layers) - # Returns distribution - def forward(self, x): - raw_init_std = np.log(np.exp(self.init_std) - 1) - x = self.model(x) - if self.dist == "tanh_normal": - mean, std = torch.chunk(x, 2, dim=-1) - mean = self.mean_scale * torch.tanh(mean / self.mean_scale) - std = self.softplus(std + raw_init_std) + self.min_std - dist = td.Normal(mean, std) - transforms = [TanhBijector()] - dist = td.transformed_distribution.TransformedDistribution( - dist, transforms) - dist = td.Independent(dist, 1) - elif self.dist == "onehot": - dist = td.OneHotCategorical(logits=x) - raise NotImplementedError("Atari not implemented yet!") - return dist + # Returns distribution + def forward(self, x): + raw_init_std = np.log(np.exp(self.init_std) - 1) + x = self.model(x) + if self.dist == "tanh_normal": + mean, std = torch.chunk(x, 2, dim=-1) + mean = self.mean_scale * torch.tanh(mean / self.mean_scale) + std = self.softplus(std + raw_init_std) + self.min_std + dist = td.Normal(mean, std) + transforms = [TanhBijector()] + dist = td.transformed_distribution.TransformedDistribution( + dist, transforms) + dist = td.Independent(dist, 1) + elif self.dist == "onehot": + dist = td.OneHotCategorical(logits=x) + raise NotImplementedError("Atari not implemented yet!") + return dist # Represents TD model in PlaNET -if torch: +class RSSM(nn.Module): + """RSSM is the core recurrent part of the PlaNET module. It consists of + two networks, one (obs) to calculate posterior beliefs and states and + the second (img) to calculate prior beliefs and states. The prior network + takes in the previous state and action, while the posterior network takes + in the previous state, action, and a latent embedding of the most recent + observation. + """ - class RSSM(nn.Module): - """RSSM is the core recurrent part of the PlaNET module. It consists of - two networks, one (obs) to calculate posterior beliefs and states and - the second (img) to calculate prior beliefs and states. The prior network - takes in the previous state and action, while the posterior network takes - in the previous state, action, and a latent embedding of the most recent - observation. - """ + def __init__(self, + action_size: int, + embed_size: int, + stoch: int = 30, + deter: int = 200, + hidden: int = 200, + act: ActFunc = None): + """Initializes RSSM - def __init__(self, - action_size: int, - embed_size: int, - stoch: int = 30, - deter: int = 200, - hidden: int = 200, - act: ActFunc = None): - """Initializes RSSM - - Args: + Args: action_size (int): Action space size embed_size (int): Size of ConvEncoder embedding stoch (int): Size of the distributional hidden state deter (int): Size of the deterministic hidden state hidden (int): General size of hidden layers act (Any): Activation function - """ - super().__init__() - self.stoch_size = stoch - self.deter_size = deter - self.hidden_size = hidden - self.act = act - if act is None: - self.act = nn.ELU + """ + super().__init__() + self.stoch_size = stoch + self.deter_size = deter + self.hidden_size = hidden + self.act = act + if act is None: + self.act = nn.ELU - self.obs1 = Linear(embed_size + deter, hidden) - self.obs2 = Linear(hidden, 2 * stoch) + self.obs1 = Linear(embed_size + deter, hidden) + self.obs2 = Linear(hidden, 2 * stoch) - self.cell = GRUCell(self.hidden_size, hidden_size=self.deter_size) - self.img1 = Linear(stoch + action_size, hidden) - self.img2 = Linear(deter, hidden) - self.img3 = Linear(hidden, 2 * stoch) + self.cell = GRUCell(self.hidden_size, hidden_size=self.deter_size) + self.img1 = Linear(stoch + action_size, hidden) + self.img2 = Linear(deter, hidden) + self.img3 = Linear(hidden, 2 * stoch) - self.softplus = nn.Softplus + self.softplus = nn.Softplus - self.device = (torch.device("cuda") if torch.cuda.is_available() - else torch.device("cpu")) + self.device = (torch.device("cuda") + if torch.cuda.is_available() else torch.device("cpu")) - def get_initial_state(self, batch_size: int) -> List[TensorType]: - """Returns the inital state for the RSSM, which consists of mean, std - for the stochastic state, the sampled stochastic hidden state - (from mean, std), and the deterministic hidden state, which is pushed - through the GRUCell. + def get_initial_state(self, batch_size: int) -> List[TensorType]: + """Returns the inital state for the RSSM, which consists of mean, + std for the stochastic state, the sampled stochastic hidden state + (from mean, std), and the deterministic hidden state, which is + pushed through the GRUCell. - Args: + Args: batch_size (int): Batch size for initial state - Returns: + Returns: List of tensors - """ - return [ - torch.zeros(batch_size, self.stoch_size).to(self.device), - torch.zeros(batch_size, self.stoch_size).to(self.device), - torch.zeros(batch_size, self.stoch_size).to(self.device), - torch.zeros(batch_size, self.deter_size).to(self.device), - ] + """ + return [ + torch.zeros(batch_size, self.stoch_size).to(self.device), + torch.zeros(batch_size, self.stoch_size).to(self.device), + torch.zeros(batch_size, self.stoch_size).to(self.device), + torch.zeros(batch_size, self.deter_size).to(self.device), + ] - def observe(self, - embed: TensorType, - action: TensorType, - state: List[TensorType] = None - ) -> Tuple[List[TensorType], List[TensorType]]: - """Returns the corresponding states from the embedding from ConvEncoder - and actions. This is accomplished by rolling out the RNN from the - starting state through eacn index of embed and action, saving all - intermediate states between. + def observe(self, + embed: TensorType, + action: TensorType, + state: List[TensorType] = None + ) -> Tuple[List[TensorType], List[TensorType]]: + """Returns the corresponding states from the embedding from ConvEncoder + and actions. This is accomplished by rolling out the RNN from the + starting state through eacn index of embed and action, saving all + intermediate states between. - Args: + Args: embed (TensorType): ConvEncoder embedding action (TensorType): Actions state (List[TensorType]): Initial state before rollout - Returns: + Returns: Posterior states and prior states (both List[TensorType]) - """ - if state is None: - state = self.get_initial_state(action.size()[0]) + """ + if state is None: + state = self.get_initial_state(action.size()[0]) - embed = embed.permute(1, 0, 2) - action = action.permute(1, 0, 2) + embed = embed.permute(1, 0, 2) + action = action.permute(1, 0, 2) - priors = [[] for i in range(len(state))] - posts = [[] for i in range(len(state))] - last = (state, state) - for index in range(len(action)): - # Tuple of post and prior - last = self.obs_step(last[0], action[index], embed[index]) - [o.append(l) for l, o in zip(last[0], posts)] - [o.append(l) for l, o in zip(last[1], priors)] + priors = [[] for i in range(len(state))] + posts = [[] for i in range(len(state))] + last = (state, state) + for index in range(len(action)): + # Tuple of post and prior + last = self.obs_step(last[0], action[index], embed[index]) + [o.append(l) for l, o in zip(last[0], posts)] + [o.append(l) for l, o in zip(last[1], priors)] - prior = [torch.stack(x, dim=0) for x in priors] - post = [torch.stack(x, dim=0) for x in posts] + prior = [torch.stack(x, dim=0) for x in priors] + post = [torch.stack(x, dim=0) for x in posts] - prior = [e.permute(1, 0, 2) for e in prior] - post = [e.permute(1, 0, 2) for e in post] + prior = [e.permute(1, 0, 2) for e in prior] + post = [e.permute(1, 0, 2) for e in post] - return post, prior + return post, prior - def imagine(self, action: TensorType, - state: List[TensorType] = None) -> List[TensorType]: - """Imagines the trajectory starting from state through a list of actions. - Similar to observe(), requires rolling out the RNN for each timestep. + def imagine(self, action: TensorType, + state: List[TensorType] = None) -> List[TensorType]: + """Imagines the trajectory starting from state through a list of actions. + Similar to observe(), requires rolling out the RNN for each timestep. - Args: + Args: action (TensorType): Actions state (List[TensorType]): Starting state before rollout - Returns: + Returns: Prior states - """ - if state is None: - state = self.get_initial_state(action.size()[0]) + """ + if state is None: + state = self.get_initial_state(action.size()[0]) - action = action.permute(1, 0, 2) + action = action.permute(1, 0, 2) - indices = range(len(action)) - priors = [[] for _ in range(len(state))] - last = state - for index in indices: - last = self.img_step(last, action[index]) - [o.append(l) for l, o in zip(last, priors)] + indices = range(len(action)) + priors = [[] for _ in range(len(state))] + last = state + for index in indices: + last = self.img_step(last, action[index]) + [o.append(l) for l, o in zip(last, priors)] - prior = [torch.stack(x, dim=0) for x in priors] - prior = [e.permute(1, 0, 2) for e in prior] - return prior + prior = [torch.stack(x, dim=0) for x in priors] + prior = [e.permute(1, 0, 2) for e in prior] + return prior - def obs_step(self, prev_state: TensorType, prev_action: TensorType, - embed: TensorType - ) -> Tuple[List[TensorType], List[TensorType]]: - """Runs through the posterior model and returns the posterior state + def obs_step( + self, prev_state: TensorType, prev_action: TensorType, + embed: TensorType) -> Tuple[List[TensorType], List[TensorType]]: + """Runs through the posterior model and returns the posterior state - Args: + Args: prev_state (TensorType): The previous state prev_action (TensorType): The previous action embed (TensorType): Embedding from ConvEncoder - Returns: + Returns: Post and Prior state - """ - prior = self.img_step(prev_state, prev_action) - x = torch.cat([prior[3], embed], dim=-1) - x = self.obs1(x) - x = self.act()(x) - x = self.obs2(x) - mean, std = torch.chunk(x, 2, dim=-1) - std = self.softplus()(std) + 0.1 - stoch = self.get_dist(mean, std).rsample() - post = [mean, std, stoch, prior[3]] - return post, prior + """ + prior = self.img_step(prev_state, prev_action) + x = torch.cat([prior[3], embed], dim=-1) + x = self.obs1(x) + x = self.act()(x) + x = self.obs2(x) + mean, std = torch.chunk(x, 2, dim=-1) + std = self.softplus()(std) + 0.1 + stoch = self.get_dist(mean, std).rsample() + post = [mean, std, stoch, prior[3]] + return post, prior - def img_step(self, prev_state: TensorType, - prev_action: TensorType) -> List[TensorType]: - """Runs through the prior model and returns the prior state + def img_step(self, prev_state: TensorType, + prev_action: TensorType) -> List[TensorType]: + """Runs through the prior model and returns the prior state - Args: + Args: prev_state (TensorType): The previous state prev_action (TensorType): The previous action - Returns: + Returns: Prior state - """ - x = torch.cat([prev_state[2], prev_action], dim=-1) - x = self.img1(x) - x = self.act()(x) - deter = self.cell(x, prev_state[3]) - x = deter - x = self.img2(x) - x = self.act()(x) - x = self.img3(x) - mean, std = torch.chunk(x, 2, dim=-1) - std = self.softplus()(std) + 0.1 - stoch = self.get_dist(mean, std).rsample() - return [mean, std, stoch, deter] + """ + x = torch.cat([prev_state[2], prev_action], dim=-1) + x = self.img1(x) + x = self.act()(x) + deter = self.cell(x, prev_state[3]) + x = deter + x = self.img2(x) + x = self.act()(x) + x = self.img3(x) + mean, std = torch.chunk(x, 2, dim=-1) + std = self.softplus()(std) + 0.1 + stoch = self.get_dist(mean, std).rsample() + return [mean, std, stoch, deter] - def get_feature(self, state: List[TensorType]) -> TensorType: - # Constructs feature for input to reward, decoder, actor, critic - return torch.cat([state[2], state[3]], dim=-1) + def get_feature(self, state: List[TensorType]) -> TensorType: + # Constructs feature for input to reward, decoder, actor, critic + return torch.cat([state[2], state[3]], dim=-1) - def get_dist(self, mean: TensorType, std: TensorType) -> TensorType: - return td.Normal(mean, std) + def get_dist(self, mean: TensorType, std: TensorType) -> TensorType: + return td.Normal(mean, std) # Represents all models in Dreamer, unifies them all into a single interface -if torch: +class DreamerModel(TorchModelV2, nn.Module): + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + super().__init__(obs_space, action_space, num_outputs, model_config, + name) - class DreamerModel(TorchModelV2, nn.Module): - def __init__(self, obs_space, action_space, num_outputs, model_config, - name): - super().__init__(obs_space, action_space, num_outputs, - model_config, name) + nn.Module.__init__(self) + self.depth = model_config["depth_size"] + self.deter_size = model_config["deter_size"] + self.stoch_size = model_config["stoch_size"] + self.hidden_size = model_config["hidden_size"] - nn.Module.__init__(self) - self.depth = model_config["depth_size"] - self.deter_size = model_config["deter_size"] - self.stoch_size = model_config["stoch_size"] - self.hidden_size = model_config["hidden_size"] + self.action_size = action_space.shape[0] - self.action_size = action_space.shape[0] + self.encoder = ConvEncoder(self.depth) + self.decoder = ConvDecoder( + self.stoch_size + self.deter_size, depth=self.depth) + self.reward = DenseDecoder(self.stoch_size + self.deter_size, 1, 2, + self.hidden_size) + self.dynamics = RSSM( + self.action_size, + 32 * self.depth, + stoch=self.stoch_size, + deter=self.deter_size) + self.actor = ActionDecoder(self.stoch_size + self.deter_size, + self.action_size, 4, self.hidden_size) + self.value = DenseDecoder(self.stoch_size + self.deter_size, 1, 3, + self.hidden_size) + self.state = None - self.encoder = ConvEncoder(self.depth) - self.decoder = ConvDecoder( - self.stoch_size + self.deter_size, depth=self.depth) - self.reward = DenseDecoder(self.stoch_size + self.deter_size, 1, 2, - self.hidden_size) - self.dynamics = RSSM( - self.action_size, - 32 * self.depth, - stoch=self.stoch_size, - deter=self.deter_size) - self.actor = ActionDecoder(self.stoch_size + self.deter_size, - self.action_size, 4, self.hidden_size) - self.value = DenseDecoder(self.stoch_size + self.deter_size, 1, 3, - self.hidden_size) - self.state = None + self.device = (torch.device("cuda") + if torch.cuda.is_available() else torch.device("cpu")) - self.device = (torch.device("cuda") if torch.cuda.is_available() - else torch.device("cpu")) + def policy(self, obs: TensorType, state: List[TensorType], explore=True + ) -> Tuple[TensorType, List[float], List[TensorType]]: + """Returns the action. Runs through the encoder, recurrent model, + and policy to obtain action. + """ + if state is None: + self.initial_state() + else: + self.state = state + post = self.state[:4] + action = self.state[4] - def policy(self, - obs: TensorType, - state: List[TensorType], - explore=True - ) -> Tuple[TensorType, List[float], List[TensorType]]: - """Returns the action. Runs through the encoder, recurrent model, - and policy to obtain action. - """ - if state is None: - self.initial_state() - else: - self.state = state - post = self.state[:4] - action = self.state[4] + embed = self.encoder(obs) + post, _ = self.dynamics.obs_step(post, action, embed) + feat = self.dynamics.get_feature(post) - embed = self.encoder(obs) - post, _ = self.dynamics.obs_step(post, action, embed) - feat = self.dynamics.get_feature(post) + action_dist = self.actor(feat) + if explore: + action = action_dist.sample() + else: + action = action_dist.mean + logp = action_dist.log_prob(action) - action_dist = self.actor(feat) - if explore: - action = action_dist.sample() - else: - action = action_dist.mean - logp = action_dist.log_prob(action) + self.state = post + [action] + return action, logp, self.state - self.state = post + [action] - return action, logp, self.state + def imagine_ahead(self, state: List[TensorType], + horizon: int) -> TensorType: + """Given a batch of states, rolls out more state of length horizon. + """ + start = [] + for s in state: + s = s.contiguous().detach() + shpe = [-1] + list(s.size())[2:] + start.append(s.view(*shpe)) - def imagine_ahead(self, state: List[TensorType], - horizon: int) -> TensorType: - """Given a batch of states, rolls out more state of length horizon. - """ - start = [] - for s in state: - s = s.contiguous().detach() - shpe = [-1] + list(s.size())[2:] - start.append(s.view(*shpe)) + def next_state(state): + feature = self.dynamics.get_feature(state).detach() + action = self.actor(feature).rsample() + next_state = self.dynamics.img_step(state, action) + return next_state - def next_state(state): - feature = self.dynamics.get_feature(state).detach() - action = self.actor(feature).rsample() - next_state = self.dynamics.img_step(state, action) - return next_state + last = start + outputs = [[] for i in range(len(start))] + for _ in range(horizon): + last = next_state(last) + [o.append(l) for l, o in zip(last, outputs)] + outputs = [torch.stack(x, dim=0) for x in outputs] - last = start - outputs = [[] for i in range(len(start))] - for _ in range(horizon): - last = next_state(last) - [o.append(l) for l, o in zip(last, outputs)] - outputs = [torch.stack(x, dim=0) for x in outputs] + imag_feat = self.dynamics.get_feature(outputs) + return imag_feat - imag_feat = self.dynamics.get_feature(outputs) - return imag_feat + def get_initial_state(self) -> List[TensorType]: + self.state = self.dynamics.get_initial_state(1) + [ + torch.zeros(1, self.action_space.shape[0]).to(self.device) + ] + return self.state - def get_initial_state(self) -> List[TensorType]: - self.state = self.dynamics.get_initial_state(1) + [ - torch.zeros(1, self.action_space.shape[0]).to(self.device) - ] - return self.state - - def value_function(self) -> TensorType: - return None + def value_function(self) -> TensorType: + return None diff --git a/rllib/agents/dreamer/utils.py b/rllib/agents/dreamer/utils.py index e86707ade..f1ba7fd70 100644 --- a/rllib/agents/dreamer/utils.py +++ b/rllib/agents/dreamer/utils.py @@ -1,6 +1,7 @@ -from ray.rllib.utils.framework import try_import_torch import numpy as np +from ray.rllib.utils.framework import try_import_torch + torch, nn = try_import_torch() # Custom initialization for different types of layers @@ -15,9 +16,6 @@ if torch: if self.bias is not None: nn.init.zeros_(self.bias) - -if torch: - class Conv2d(nn.Conv2d): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -27,9 +25,6 @@ if torch: if self.bias is not None: nn.init.zeros_(self.bias) - -if torch: - class ConvTranspose2d(nn.ConvTranspose2d): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -39,9 +34,6 @@ if torch: if self.bias is not None: nn.init.zeros_(self.bias) - -if torch: - class GRUCell(nn.GRUCell): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -52,10 +44,7 @@ if torch: nn.init.zeros_(self.bias_ih) nn.init.zeros_(self.bias_hh) - -# Custom Tanh Bijector due to big gradients through Dreamer Actor -if torch: - + # Custom Tanh Bijector due to big gradients through Dreamer Actor class TanhBijector(torch.distributions.Transform): def __init__(self): super().__init__() diff --git a/rllib/models/tests/test_attention_nets.py b/rllib/models/tests/test_attention_nets.py index f8b8e4be2..ac6ec134d 100644 --- a/rllib/models/tests/test_attention_nets.py +++ b/rllib/models/tests/test_attention_nets.py @@ -16,7 +16,7 @@ torch, nn = try_import_torch() tf1, tf, tfv = try_import_tf() -class TestModules(unittest.TestCase): +class TestAttentionNets(unittest.TestCase): """Tests various torch/modules and tf/layers required for AttentionNet""" def train_torch_full_model(self, diff --git a/rllib/models/tests/test_convtranspose2d_stack.py b/rllib/models/tests/test_convtranspose2d_stack.py new file mode 100644 index 000000000..26f50a28c --- /dev/null +++ b/rllib/models/tests/test_convtranspose2d_stack.py @@ -0,0 +1,64 @@ +import cv2 +import gym +import numpy as np +import os +from pathlib import Path +import unittest + +from ray.rllib.models.preprocessors import GenericPixelPreprocessor +from ray.rllib.models.torch.modules.convtranspose2d_stack import \ + ConvTranspose2DStack +from ray.rllib.utils.framework import try_import_torch, try_import_tf + +torch, nn = try_import_torch() +tf1, tf, tfv = try_import_tf() + + +class TestConvTranspose2DStack(unittest.TestCase): + """Tests our ConvTranspose2D Stack modules/layers.""" + + def test_convtranspose2d_stack(self): + """Tests, whether the conv2d stack can be trained to predict an image. + """ + batch_size = 128 + input_size = 1 + module = ConvTranspose2DStack(input_size=input_size) + preprocessor = GenericPixelPreprocessor( + gym.spaces.Box(0, 255, (64, 64, 3), np.uint8), options={"dim": 64}) + optim = torch.optim.Adam(module.parameters(), lr=0.0001) + + rllib_dir = Path(__file__).parent.parent.parent + img_file = os.path.join(rllib_dir, + "tests/data/images/obstacle_tower.png") + img = cv2.imread(img_file).astype(np.float32) + # Preprocess. + img = preprocessor.transform(img) + # Make channels first. + img = np.transpose(img, (2, 0, 1)) + # Add batch rank and repeat. + imgs = np.reshape(img, (1, ) + img.shape) + imgs = np.repeat(imgs, batch_size, axis=0) + # Move to torch. + imgs = torch.from_numpy(imgs) + init_loss = loss = None + for _ in range(10): + # Random inputs. + inputs = torch.from_numpy( + np.random.normal(0.0, 1.0, (batch_size, input_size))).float() + distribution = module(inputs) + # Construct a loss. + loss = -torch.mean(distribution.log_prob(imgs)) + if init_loss is None: + init_loss = loss + print("loss={}".format(loss)) + # Minimize loss. + loss.backward() + optim.step() + self.assertLess(loss, init_loss) + + +if __name__ == "__main__": + import pytest + import sys + + sys.exit(pytest.main(["-v", __file__])) diff --git a/rllib/models/tf/layers/noisy_layer.py b/rllib/models/tf/layers/noisy_layer.py index 6e6a32c68..f7945d3f8 100644 --- a/rllib/models/tf/layers/noisy_layer.py +++ b/rllib/models/tf/layers/noisy_layer.py @@ -7,13 +7,15 @@ tf1, tf, tfv = try_import_tf() class NoisyLayer(tf.keras.layers.Layer if tf else object): - """A Layer that adds learnable Noise - a common dense layer: y = w^{T}x + b - a noisy layer: y = (w + \\epsilon_w*\\sigma_w)^{T}x + + """A Layer that adds learnable Noise to some previous layer's outputs. + + Consists of: + - a common dense layer: y = w^{T}x + b + - a noisy layer: y = (w + \\epsilon_w*\\sigma_w)^{T}x + (b+\\epsilon_b*\\sigma_b) - where \epsilon are random variables sampled from factorized normal + , where \epsilon are random variables sampled from factorized normal distributions and \\sigma are trainable variables which are expected to - vanish along the training procedure + vanish along the training procedure. """ def __init__(self, prefix, out_size, sigma0, activation="relu"): diff --git a/rllib/models/torch/misc.py b/rllib/models/torch/misc.py index b7f352434..977a5d907 100644 --- a/rllib/models/torch/misc.py +++ b/rllib/models/torch/misc.py @@ -1,5 +1,6 @@ """ Code adapted from https://github.com/ikostrikov/pytorch-a3c""" import numpy as np +from typing import List from ray.rllib.utils.framework import get_activation_fn, try_import_torch @@ -138,3 +139,15 @@ class AppendBiasLayer(nn.Module): out = torch.cat( [x, self.log_std.unsqueeze(0).repeat([len(x), 1])], axis=1) return out + + +class Reshape(nn.Module): + """Standard module that reshapes/views a tensor + """ + + def __init__(self, shape: List): + super().__init__() + self.shape = shape + + def forward(self, x): + return x.view(*self.shape) diff --git a/rllib/models/torch/modules/convtranspose2d_stack.py b/rllib/models/torch/modules/convtranspose2d_stack.py new file mode 100644 index 000000000..0eab6dc7c --- /dev/null +++ b/rllib/models/torch/modules/convtranspose2d_stack.py @@ -0,0 +1,77 @@ +from typing import Tuple + +from ray.rllib.models.torch.misc import Reshape +from ray.rllib.models.utils import get_initializer +from ray.rllib.utils.framework import get_activation_fn, try_import_torch + +torch, nn = try_import_torch() +if torch: + import torch.distributions as td + + +class ConvTranspose2DStack(nn.Module): + """ConvTranspose2D decoder generating an image distribution from a vector. + """ + + def __init__(self, + *, + input_size: int, + filters: Tuple[Tuple[int]] = ((1024, 5, 2), (128, 5, 2), + (64, 6, 2), (32, 6, 2)), + initializer="default", + bias_init=0, + activation_fn: str = "relu", + output_shape: Tuple[int] = (3, 64, 64)): + """Initializes a TransposedConv2DStack instance. + + Args: + input_size (int): The size of the 1D input vector, from which to + generate the image distribution. + filters (Tuple[Tuple[int]]): Tuple of filter setups (1 for each + ConvTranspose2D layer): [in_channels, kernel, stride]. + initializer (Union[str]): + bias_init (float): The initial bias values to use. + activation_fn (str): Activation function descriptor (str). + output_shape (Tuple[int]): Shape of the final output image. + """ + super().__init__() + self.activation = get_activation_fn(activation_fn, framework="torch") + self.output_shape = output_shape + initializer = get_initializer(initializer, framework="torch") + + in_channels = filters[0][0] + self.layers = [ + # Map from 1D-input vector to correct initial size for the + # Conv2DTransposed stack. + nn.Linear(input_size, in_channels), + # Reshape from the incoming 1D vector (input_size) to 1x1 image + # format (channels first). + Reshape([-1, in_channels, 1, 1]), + ] + for i, (_, kernel, stride) in enumerate(filters): + out_channels = filters[i + 1][0] if i < len(filters) - 1 else \ + output_shape[0] + conv_transp = nn.ConvTranspose2d(in_channels, out_channels, kernel, + stride) + # Apply initializer. + initializer(conv_transp.weight) + nn.init.constant_(conv_transp.bias, bias_init) + self.layers.append(conv_transp) + # Apply activation function, if provided and if not last layer. + if self.activation is not None and i < len(filters) - 1: + self.layers.append(self.activation()) + + # num-outputs == num-inputs for next layer. + in_channels = out_channels + + self._model = nn.Sequential(*self.layers) + + def forward(self, x): + # x is [batch, hor_length, input_size] + batch_dims = x.shape[:-1] + model_out = self._model(x) + + # Equivalent to making a multivariate diag. + reshape_size = batch_dims + self.output_shape + mean = model_out.view(*reshape_size) + return td.Independent(td.Normal(mean, 1.0), len(self.output_shape)) diff --git a/rllib/models/torch/modules/noisy_layer.py b/rllib/models/torch/modules/noisy_layer.py index 67d92fc15..3e835b66b 100644 --- a/rllib/models/torch/modules/noisy_layer.py +++ b/rllib/models/torch/modules/noisy_layer.py @@ -7,13 +7,15 @@ torch, nn = try_import_torch() class NoisyLayer(nn.Module): - """A Layer that adds learnable Noise - a common dense layer: y = w^{T}x + b - a noisy layer: y = (w + \\epsilon_w*\\sigma_w)^{T}x + + """A Layer that adds learnable Noise to some previous layer's outputs. + + Consists of: + - a common dense layer: y = w^{T}x + b + - a noisy layer: y = (w + \\epsilon_w*\\sigma_w)^{T}x + (b+\\epsilon_b*\\sigma_b) - where \epsilon are random variables sampled from factorized normal + , where \epsilon are random variables sampled from factorized normal distributions and \\sigma are trainable variables which are expected to - vanish along the training procedure + vanish along the training procedure. """ def __init__(self, in_size, out_size, sigma0, activation="relu"): diff --git a/rllib/models/utils.py b/rllib/models/utils.py index a18c01676..2c9f076f0 100644 --- a/rllib/models/utils.py +++ b/rllib/models/utils.py @@ -1,3 +1,6 @@ +from ray.rllib.utils.framework import try_import_tf, try_import_torch + + def get_filter_config(shape): """Returns a default Conv2D filter config (list) for a given image shape. @@ -30,3 +33,35 @@ def get_filter_config(shape): "Default configurations are only available for inputs of shape " "[42, 42, K] and [84, 84, K]. You may alternatively want " "to use a custom model or preprocessor.") + + +def get_initializer(name, framework="tf"): + """Returns a framework specific initializer, given a name string. + + Args: + name (str): One of "xavier_uniform" (default), "xavier_normal". + framework (str): One of "tf" or "torch". + + Returns: + A framework-specific initializer function, e.g. + tf.keras.initializers.GlorotUniform or + torch.nn.init.xavier_uniform_. + + Raises: + ValueError: If name is an unknown initializer. + """ + if framework == "torch": + _, nn = try_import_torch() + if name in [None, "default", "xavier_uniform"]: + return nn.init.xavier_uniform_ + elif name == "xavier_normal": + return nn.init.xavier_normal_ + else: + tf1, tf, tfv = try_import_tf() + if name in [None, "default", "xavier_uniform"]: + return tf.keras.initializers.GlorotUniform + elif name == "xavier_normal": + return tf.keras.initializers.GlorotNormal + + raise ValueError("Unknown activation ({}) for framework={}!".format( + name, framework)) diff --git a/rllib/tests/data/images/obstacle_tower.png b/rllib/tests/data/images/obstacle_tower.png new file mode 100644 index 0000000000000000000000000000000000000000..0b2231ad6804ff78767fa7d9a56d34bd283c1179 GIT binary patch literal 8737 zcmV++BHrDJP)00001b5ch_0Itp) z=>Px#1ZP1_K>z@;j|==^1poj532;bRa{vGi!vFvd!vV){sAK>DA+||GK~!i%O?nBi z9ani~pL4oTuXla#-M8xLNwO?~ZEUQuGY^QYmCJErV%gNnwzj zBrygpR1FAGedL&ZqC=kNILE+6=th(9N}!^^N8CJ77{|#+QJT477TQp%fR#+s#C1$! zOAC`sw8Cw;g7e7Wu50bvJ6Ch1v@|RX>3!D|dFn~g9Sox+Gq|9|F$)*0=JJJUzHiy` zB)5c?X6bnzrc;SBd1+e=LPLy+i!eCH@}gwWnf@?b+Z&|jg?=(lbY7B@rQ&&xI}K@Z zQ`m-4)~u58vgTQqEeByFOfrYc3zHSl2Z}kK8#p)>b9L|RTG$AzB2P_N6zXApVKQN0 zO86f}UaBN&t5!3or`1m@C>z|RWpJ8Ztj!YHYNoGYc^5Ddfy?wmZx zja=mM!c=_}XP2X{R1;(221Twz(26uwgfYdgsi+E7^-D6yw@+B|o~@0t z%0g1L(8x9;OH+6SMt7xfOo&w=?n8x`%RNUzJ`+J|m6)P5DEp;p-1jX-W<^vm!4umB&e;6xKw@ zO3n4DGux-0OFWU>^~=re?U`siGA*f@QTG^xLH?zsF-{arg}{I!NGQk)msFTW0ELL- z@C`z@Gy+@02QWToCc@Ve0)mtBrUeJ4}xt{xAxBDTG;;5KJLb^OD)H zgkdFFYGiq;U@>x;fk{%uRBo09pa2WMh81<&_9E84|ABBbIsZTZ*TN@1K3(bzU{|Zz zwyn%Et2q`PpTr3yu2SOifP5yfkzfK5$@`%JOqqOzu))A1xZf70rbwa4g1Y!qX$~`_ zX=)3a1kh9>NDKllq`2*v$Xm#LaeIRux%P^{hD(iprUp?7*CS4Jq4v%M6QINzmza=HrN6wTGry8V^54Gu%YePnm?KrfDAtB3bJC zB{=VOL$P;u%0fbwMsK8asZpO0zX;84REBBM598^k5BW_)!JuW4b73Lhjz6_~><7m; zH(ok*3OJZ2Nnu({Wkvz0MW#b~Oap+e2ALV3O9kWxyeTc}mUdltm=?n@@;oF#krrAq z13`N9{ylR*55#UiOmL*7 ztAfHAKnJN|ZS{<hUFh({S-EK@I{v&4&?NT^;ZCRDp#{z!7Q?@E3?kPJ<1o(pIzK zBgBH&5n!m(V-)2wgO!lsuEmED>P8_gKpi&(NdnKKP6EPxNgRP$?Bt6VyW?2r&g-wb za%IMh2i-VTz8^&6RP3AY0A;uPv0^5o2S5u{!R$Ho%_%V%1NOnFY%t7+35zlv79vbV zyB1iMIT)1Oa+BQ15|ndT4iiv(8yr~#dO>jwb0D$t5xP)G0`f?LCuFW;+bF)$b{wHW zHw>s)GC*gX@wvs>EE&T!DlHb~nvv4&x^3wQO8{@iSyA6ABo8xk^AQXd}1= z?%-<#F{VXa_=tzu7IynIaL7E}vDdkxXYJgw>bV=zL<=A$zg?6pea&1=v7_=}G?(Rgz zAWD%^^f*kD^Abe}45MVAtGI;6X>|}P2UUhkCg6Hj2s)r1al>=%r6%0P09lnA848)} zH&8WDWMFwzBFqdsC`dzv2Y7(hF(s9n%OM11CP!F;=h}9?i`x-I-0HPea8 zcCzPs@4uSDM{>fNK!1EGRu(!Y z7?_-dun}g;?A~1e%%XhXi)YX7<8S#1`y&@sqAVQSp~B{Q9Gyu~P(cM^A!AXi2!Kfk z`h(DwGMCd}G7zs2aiR|@5w755qH&z0fB}5OIZkjVAyRb%6@8uaf}v-rs!iEXjqOk& z`{M%b1fqb(+4=zkhLhbEl85@OYY`)-%uq&>Khx zjsr=+>M$Fy2kt`9;4(wgRDfw2U@`&H#3oy|KvoJ4AT9;s&qHOfO;T*yciwsKOLi8C zj(6?+p=7pocjLD2cE4bIje4y%-Ez9UQIcljiajlq5X(j@lh1YBJ)IHYEO0!}aZ$mu zq@wFpO+tjiHCaximkYPo>+bEe!5l|`)-q4EKmbq!3baJ_z=UxQo2Vj3xj+Mj9s)oB zzKGn0!va?#`q5rhU_g}`#+k8lduyecE#FuC$34SaU&)_+MZTJDhHInIK$yes+4IfW z{dqTcQeu!)p|m)-uZ8Gn)O54S=Vvbc-XDDC*iRoSlo!P*6!Rn-qd+0k38E-S$ix6X z!+FEO>0HIVq*&-ofdh6rfi6<6nKHy+n0AxG5gay-D7C08Qb94{4Kf7a9|}MbM^|{YqckjP@$0Ofg|EGt4cGFSupFjPMAN}wQ>}K zu_wt56{#aISPt$6f&y^SumplqMKD;2b^z_HfszUuHzS%1%g%2Gup0KSxz!>F$THq=RDwH&M(jlf>-5XVPraHz6c`MZ0EGdfAcBlCyZ&(Jr>8GXIXh208xKcdI9M18LrwTVUbr29K=!GL3;S1^>dIKAa>V((|tEfyshQ8QS{bC*Hnz@Stn= zKl$lAGgd=RwH|r=IU_q)COemQ&u+vUQ6Bw$@U7SH2;tU5iI@Vs zAxNE5G}pzrhufqasq18*oHk zXv+p+{gXf3v@DVy*TJB0ndGL*K?4tBwJf8eGsz7cCl@E*PL7YAji5=@XxBh14T<*1x8QT_QEr-OHPWf6E7#C{o#m#NLwZ001sq)AKW*S_ASK7??_L$59)rEynW^hDAzyw- z@tUh^yIEPdvIfd#3P}_dp@jkrhKid(3v&Uy!$mX(HY=o&S1gzY{0Z6*y=^3b4O+(P zu4Ss3%Zf_cveqsBnoU!ejIZ@Ena%q2ok>eKmbC)AaaGcLN5C(;ZA# zF*$e?hz^ERp$bHqfXf71v_>n9>C(7=)zZw&j6Vzu(`kjJG1Ck{4=G`UOjOy6Y2O!f z0)<@QF>N-*n_V-+CHZ7e>V}0{T%3 zOk7b#GwfO}b^Iia-*R*CqlZT@&~O+dt4Krmj%&YWhvOLcRzO5V0oDWlgN|Y@6bm2{ zG6#5tKif2=Xq8K?ca6po9}toTi^NXx4iehn_P#X#~-7Rg4%-JfRHipP*3~`(eKu54&hSRY@F|N1^$#k39H+-~Q%@{=*MG_VGtQ z{h253`SmYu4UJ3P{8B$9<5#0{KG~DF5QscbK|xNRA-dtt>eT_IdBuluBi*nDN71QyfbS2LZWvrAM{cTXcz3!Ot%v$4NGUusyt`R;ep zo(cFYW->s_NfK!>v9j1&-$ob0(~O}pjAGb2pBXVw`1BRwjy#!U!T2np%Y@RdHT|ZUCWyq`9JPtf&NDQD2#4MXJyx z056QJ(7*hK=zEVnzJBW2n{RyE?zo$q?z2z)OdzDSea8)l9{b73h)N$X(16np3tGy> zc5j5oHKG9F3n<44EK<^&CD34egE~sqh446mQBbg8>*<;a@3B$PK&5a8CBk%u>?;LyFx7fz|_Q6AYEx;S)78Mk9p`Z+?@U;0N zg=26N0R|Z+#;0#7;;4?IHi7U_V4x@&urh5T(Pj|-3iM?Tf`s-4sVd_OxQu3Sk)-&) z)UoDKZV)Lb@Xv-7E5}l0MPDg`>}wBgkGV#nK+lBv9l)%^I?d@K3kSU|o$z6t=T@+@ zxijj9U_Y=e`YbfqurxlP&>}z)=a<7~GH3y4OUXwkCTHO|!TnGS&Vk1&E6{ZtN}3=G zv@ru-H4zRvAO&1$;}gLY%?Kux!x13e{35Mqo9I&$78m>J&-#O z{_78av8PgKk9u!A9*j=2V~W{^U(pz(Y_(}ctfw*^$x0A3i61OeNkiwTCrdMp`S}{| zsn_e;xEyJu3CIUcrU87o0QW$h2~6dcVJ@o#B;lrG6Ujv%1hEvo!xTh~xqNDCD^Fs4 zyn6ou`|uPypyDX_Z|_nmIXrhbms%l}QBuUd#W}E>FhLyC%o90i{mM4gClZSm2D1T` zBb{M{QdVLhnqed)1G)uVN47(G6Q!CE6)cIP7+R(PSpb8CX>24KnZ%e-S4dM;Y8r8# zS;F)j?^D0~JA0DCo4xAFxhP7;k5;n{nGDe#orr0F=sbjdu4Z!zxB2iOh4?V zH(s*)(gDG5J(foURmD)Uj_u%{fgN1oezKF-+beN zuZ;UUqeMF?6fa;G6q3MXo}~{=$w$t%m065ZL2rEj<#mV9UH_|luHgkshsiT9Y&V-V z*DjjPI+z}uPzQ;kj<$T|0{B+EH0N_L!7M?$!oAvyr_LWbxE#2upuPn~#m@2sSc(F? zzP3Z08EY7&FKKfd9C%@AetlydD2Ql`MJ%Z((L2E`8z7#SnYOEA5G1k zeeBENK8P6)O&Tw14xbO?oTt9>?9ALsg9sAn{XOrv(Kp>P*QRu9HY!=+SgsIpfP`|< z<_j>Co!!yGwB2mjDEmp8m?$&S%2cw@sn-Li(QYm-9on<+;GX4$BKLc}@WOi6lw3M4 z;3Y}fxy@d!(XeaoIP6ymkTsf9+-*PpY_ri4 zpaD_^8UiFMK9cpRmU+W)w05cAn(EYTZK&8|)8{t(ci;P_?U!EujXRHi?_A$y#W(-# z+sN!o%o->~I{{MMas8sfD0aAXkyd=S8)osSU6YHQ#_Y7;YSw6Tx3qercoeG(TRUI> z=HvH&9}zxe(mC$$U`WyUGUUW$mxG*8Oh^UZw=cB95q!)?{=dSLlAvsHK7 z4QibzwlK;9{Tj1SF`$yPP3mT{!bB5|HZ~OjOATdOm~DF}_RP{N%PV}z3Xsg=icSj> z19yPYb;mJn3m(a$F+f}FTj;Fr8c5Pqqg5;W(HK0MUS|<{TaIn{El(n(Jkeda z1~vZZQyUAjGgE=!!1|J*4p5pQ5#0>Xp2SKRa&>*@`BUBFXSXj6iVK%w3yG0Lw#zL^ zwmkOqMZaNPxp#(MyyuxMv@O-f`tA>ZdL}FFJKph@Z+z}67f!siuEoBCGtE8X^!YK` zoTd%5GK;ZosYb2o)oP|Vdi@oq(0;vj&A~RjJJoIm(8JKtI66)q5$(RqiX<>4h zf0U@5IN2PH(%A6Fl6~rfzrs|uGb+1-XtVioN8qmFu1B*Qd&fBs_F*9j1~j~Gfa32zM_7k z8fblR1pO(zi33M+sGlj{mM-MwzVDP-= zxl8lbORo-R=G?fOdA>8%lCJC6w!{ENP&03$3V{hTOCqSGWg3=OFm7kM*xU`zUl^Wz zIl1~!1JXlKhzreiNk4lsX_J(Y6r!WmKmi5h!XZdX+NYx}M-0}nrcUi{fd$-nb!%bT zDuuJzOVDGW6l`{rf@O8PxfnEH09%-7tI;-!JKl5++QwYvj%@;YUcK1e*fEzDo4$j- z#73pBENIYv;e=@#T1p$0xw+F4zCr9jbpXT6t)E=oLh#&jWX{rP05~*2kA^{Q z6A^lL`x}4ts}DccvaEKHh_?Ur2T%G%9_!@r%9NyE^tkJT{@S_z4cF{T!XXp~RiRgT zCcW@AEkoRP^r{#gRz`d1a5<{4BkiCOoIKtC`HAySp6vIAeHACXOsDG&X@c;8TC%`n zK>$7qWvaH_K&M^8TMgf9)g9MCC&L_Rp4;s96>ry?YwJMkd~sn8^;>WW+5=C3xEFaj z(lZy<*TO=#+lycS-`_EPP-}7Xk=eRu18b$w(>1ZO+yQ^PloTgU?|6Y5rD>R@jL4@D zzjX8vV}(yY4&*+%d=~q@<5+e$OvXuJ`_6L1LJ!VS$U&#MZUzpN_gxzlACa6D3e?E7 zBM90PXqvPkptdyo8R zJVMu~E{vjR(9_c8rL}+22OOj8Pf{aOw1Efz1y;U1>(92^o4w)M#ZkLG6(>3S$OrB@ zb8gu7vpowfDYGo5Sciik48p*%_(M-`PfspB#+`qoxRBB znt*1JmHN;^Qkj=Fu{ERpZA_BVD_oR6P`=V}B%(1@k+|pftM<*eW*d%lMM*E|K((nQ zgAQo#g=wM0T-fTZ?Upr*WeVwUMaKBqQ{gKc3Buxr*R@~2x4km&xVA+-OisThf_1Sj zDjU`U3=qcyiKk#I;5HyV!nNjluGFW`ZYgHY)osW3T{ph_-VeNR;$`2^L7?b8Ltw;gBkU4XE?k)ez+=cS&^46(qL``I8}$OrsohyT zzd3y2-0m>cR~-nJrkqYIpw6r0nWFcq0_g#;1Yfw^ks`B+2;mH!CQ=$=5CXbbo}uyd zT+gu5 z><|CbuK|v66&1*Fh_}E;iuT@0I1Y;7Hhhcj>&aI(dOiKCZ@V^)dzAKR4u^n6QqLnq z;F9fO7~EfZnRJ*OF)#xLPT&Oi2M{zEjcFlO;162b5~ZLBBlg7eTZa#J_H{robe2*0 z5rbmY&#!21Opa?LB+4iO;V&J%5;g=Z0ECb?$Yo><;0;g+X(*cLF`$gKcgnMD<7dY& zUb%lZ9LGul(Lw69v@GgG5$j<5l7a^VOTMNC#KlL5jtQYE+=I*Lk*Ng3MUaDo(jw3i z+1C6tYGQPDBb=$%(65o`D#KI=BxDxS6Kj=|4d<%u6>op@wblP6{Q#V7iJR%CkoZ(d z2EDK~33RSA>n0YcFH>Sl20X2dO zrlWSLx~~M6E2D_s_s3~sOAqVYrupjH;M8h&rqy(X4eOu@flyIQKT5Ue?G`l)S*qBd ze)_!y{qBjvph~zZ$b{XhBy{O{-o6DVy^ zfPS^OJxVWbkG1IpPOv=V99Z;EpYNW%*qxqk_riF6bC~2suOI$DfjMC+E{!xn00000 LNkvXXu0mjf%Z|Eq literal 0 HcmV?d00001 diff --git a/rllib/utils/framework.py b/rllib/utils/framework.py index 3423142fe..156539178 100644 --- a/rllib/utils/framework.py +++ b/rllib/utils/framework.py @@ -219,11 +219,12 @@ def get_variable(value, return value +# TODO: (sven) move to models/utils.py def get_activation_fn(name, framework="tf"): """Returns a framework specific activation function, given a name string. Args: - name (str): One of "relu" (default), "tanh", or "linear". + name (str): One of "relu" (default), "tanh", "swish", or "linear". framework (str): One of "tf" or "torch". Returns: