def build_q_net(name_): activation = get_activation_fn( critic_hidden_activation, framework="torch") # For continuous actions: Feed obs and actions (concatenated) # through the NN. For discrete actions, only obs. q_net = nn.Sequential() ins = self.obs_ins + self.action_dim for i, n in enumerate(critic_hiddens): q_net.add_module( "{}_hidden_{}".format(name_, i), SlimFC( ins, n, initializer=torch.nn.init.xavier_uniform_, activation_fn=activation)) ins = n q_net.add_module( "{}_out".format(name_), SlimFC( ins, 1, initializer=torch.nn.init.xavier_uniform_, activation_fn=None)) return q_net
def __init__( self, obs_space: gym.spaces.Space, fcnet_hiddens_per_candidate=(256, 32), ): """Initializes a QValueModel instance. Each document candidate receives one full Q-value stack, defined by `fcnet_hiddens_per_candidate`. The input to each of these Q-value stacks is always {[user] concat [document[i]] for i in document_candidates}. Extra model kwargs: fcnet_hiddens_per_candidate: List of layer-sizes for each(!) of the candidate documents. """ super().__init__() self.orig_obs_space = obs_space self.embedding_size = self.orig_obs_space["doc"]["0"].shape[0] self.num_candidates = len(self.orig_obs_space["doc"]) assert self.orig_obs_space["user"].shape[0] == self.embedding_size self.q_nets = nn.ModuleList() for i in range(self.num_candidates): layers = nn.Sequential() ins = 2 * self.embedding_size for j, h in enumerate(fcnet_hiddens_per_candidate): layers.add_module( f"q_layer_{i}_{j}", SlimFC(in_size=ins, out_size=h, activation_fn="relu"), ) ins = h layers.add_module(f"q_out_{i}", SlimFC(ins, 1, activation_fn=None)) self.q_nets.append(layers)
def __init__(self, obs_space, action_space, num_outputs, model_config, name): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) hiddens = model_config.get("fcnet_hiddens") activation = _get_activation_fn(model_config.get("fcnet_activation")) logger.debug("Constructing fcnet {} {}".format(hiddens, activation)) layers = [] last_layer_size = np.product(obs_space.shape) for size in hiddens: layers.append( SlimFC(in_size=last_layer_size, out_size=size, initializer=normc_initializer(1.0), activation_fn=activation)) last_layer_size = size self._hidden_layers = nn.Sequential(*layers) self._logits = SlimFC(in_size=last_layer_size, out_size=num_outputs, initializer=normc_initializer(0.01), activation_fn=None) self._value_branch = SlimFC(in_size=last_layer_size, out_size=1, initializer=normc_initializer(1.0), activation_fn=None) self._cur_value = None
def __init__( self, observation_space, action_space, num_outputs, model_config, name ): TorchModelV2.__init__( self, observation_space, action_space, num_outputs, model_config, name ) nn.Module.__init__(self) # Non-shared initial layer. self.first_layer = SlimFC( int(np.product(observation_space.shape)), 64, activation_fn=nn.ReLU, initializer=torch.nn.init.xavier_uniform_, ) # Non-shared final layer. self.last_layer = SlimFC( 64, self.num_outputs, activation_fn=None, initializer=torch.nn.init.xavier_uniform_, ) self.vf = SlimFC( 64, 1, activation_fn=None, initializer=torch.nn.init.xavier_uniform_, ) self._global_shared_layer = TORCH_GLOBAL_SHARED_LAYER self._output = None
def __init__(self, obs_space, action_space, num_outputs, model_config, name): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) filters = model_config.get("conv_filters") if not filters: filters = _get_filter_config(obs_space.shape) layers = [] (w, h, in_channels) = obs_space.shape in_size = [w, h] for out_channels, kernel, stride in filters[:-1]: padding, out_size = valid_padding(in_size, kernel, [stride, stride]) layers.append( SlimConv2d(in_channels, out_channels, kernel, stride, padding)) in_channels = out_channels in_size = out_size out_channels, kernel, stride = filters[-1] layers.append( SlimConv2d(in_channels, out_channels, kernel, stride, None)) self._convs = nn.Sequential(*layers) self._logits = SlimFC( out_channels, num_outputs, initializer=nn.init.xavier_uniform_) self._value_branch = SlimFC( out_channels, 1, initializer=normc_initializer()) self._cur_value = None
def __init__(self, obs_space, action_space, num_outputs, model_config, name, cnn_shape): super().__init__(obs_space, action_space, num_outputs, model_config, name) self.lstm_state_size = 16 self.cnn_shape = list(cnn_shape) self.visual_size_in = cnn_shape[0] * cnn_shape[1] * cnn_shape[2] # MobileNetV2 has a flat output of (1000,). self.visual_size_out = 1000 # Load the MobileNetV2 from torch.hub. self.cnn_model = torch.hub.load("pytorch/vision:v0.6.0", "mobilenet_v2", pretrained=True) self.lstm = nn.LSTM(self.visual_size_out, self.lstm_state_size, batch_first=True) # Postprocess LSTM output with another hidden layer and compute values. self.logits = SlimFC(self.lstm_state_size, self.num_outputs) self.value_branch = SlimFC(self.lstm_state_size, 1) # Holds the current "base" output (before logits layer). self._features = None
def __init__(self, in_dim: int, out_dim: int, num_heads: int, head_dim: int, input_layernorm: bool = False, output_activation: Union[str, callable] = None, **kwargs): """Initializes a RelativeMultiHeadAttention nn.Module object. Args: in_dim (int): out_dim (int): The output dimension of this module. Also known as "attention dim". num_heads (int): The number of attention heads to use. Denoted `H` in [2]. head_dim (int): The dimension of a single(!) attention head Denoted `D` in [2]. input_layernorm (bool): Whether to prepend a LayerNorm before everything else. Should be True for building a GTrXL. output_activation (Union[str, callable]): Optional activation function or activation function specifier (str). Should be "relu" for GTrXL. **kwargs: """ super().__init__(**kwargs) # No bias or non-linearity. self._num_heads = num_heads self._head_dim = head_dim # 3=Query, key, and value inputs. self._qkv_layer = SlimFC(in_size=in_dim, out_size=3 * num_heads * head_dim, use_bias=False) self._linear_layer = SlimFC( in_size=num_heads * head_dim, out_size=out_dim, use_bias=False, activation_fn=output_activation, ) self._uvar = nn.Parameter(torch.zeros(num_heads, head_dim)) self._vvar = nn.Parameter(torch.zeros(num_heads, head_dim)) nn.init.xavier_uniform_(self._uvar) nn.init.xavier_uniform_(self._vvar) self.register_parameter("_uvar", self._uvar) self.register_parameter("_vvar", self._vvar) self._pos_proj = SlimFC(in_size=in_dim, out_size=num_heads * head_dim, use_bias=False) self._rel_pos_embedding = RelativePositionEmbedding(out_dim) self._input_layernorm = None if input_layernorm: self._input_layernorm = torch.nn.LayerNorm(in_dim)
def __init__(self, obs_space, action_space, num_outputs, model_config, name): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) if action_space != Tuple([Discrete(2), Discrete(2)]): raise ValueError( "This model only supports the [2, 2] action space") # Output of the model (normally 'logits', but for an autoregressive # dist this is more like a context/feature layer encoding the obs) self.context_layer = SlimFC( in_size=obs_space.shape[0], out_size=num_outputs, initializer=normc_init_torch(1.0), activation_fn=nn.Tanh, ) # V(s) self.value_branch = SlimFC( in_size=num_outputs, out_size=1, initializer=normc_init_torch(0.01), activation_fn=None, ) # P(a1 | obs) self.a1_logits = SlimFC(in_size=num_outputs, out_size=2, activation_fn=None, initializer=normc_init_torch(0.01)) class _ActionModel(nn.Module): def __init__(self): nn.Module.__init__(self) self.a2_hidden = SlimFC(in_size=1, out_size=16, activation_fn=nn.Tanh, initializer=normc_init_torch(1.0)) self.a2_logits = SlimFC(in_size=16, out_size=2, activation_fn=None, initializer=normc_init_torch(0.01)) def forward(self_, ctx_input, a1_input): a1_logits = self.a1_logits(ctx_input) a2_logits = self_.a2_logits(self_.a2_hidden(a1_input)) return a1_logits, a2_logits # P(a2 | a1) # --note: typically you'd want to implement P(a2 | a1, obs) as follows: # a2_context = tf.keras.layers.Concatenate(axis=1)( # [ctx_input, a1_input]) self.action_module = _ActionModel() self._context = None
def __init__(self, obs_space: gym.spaces.Space, action_space: gym.spaces.Space, num_outputs: int, model_config: ModelConfigDict, name: str): nn.Module.__init__(self) super().__init__(obs_space, action_space, None, model_config, name) if isinstance(action_space, Discrete): self.action_dim = action_space.n elif isinstance(action_space, MultiDiscrete): self.action_dim = np.product(action_space.nvec) elif action_space.shape is not None: self.action_dim = int(np.product(action_space.shape)) else: self.action_dim = int(len(action_space)) cfg = model_config self.attention_dim = cfg["attention_dim"] # Construct GTrXL sub-module w/ num_outputs=None (so it does not # create a logits/value output; we'll do this ourselves in this wrapper # here). self.gtrxl = GTrXLNet( obs_space, action_space, None, model_config, "gtrxl", num_transformer_units=cfg["attention_num_transformer_units"], attention_dim=self.attention_dim, num_heads=cfg["attention_num_heads"], head_dim=cfg["attention_head_dim"], memory_inference=cfg["attention_memory_inference"], memory_training=cfg["attention_memory_training"], position_wise_mlp_dim=cfg["attention_position_wise_mlp_dim"], init_gru_gate_bias=cfg["attention_init_gru_gate_bias"], ) # Set final num_outputs to correct value (depending on action space). self.num_outputs = num_outputs # Postprocess GTrXL output with another hidden layer and compute # values. self._logits_branch = SlimFC( in_size=self.attention_dim, out_size=self.num_outputs, activation_fn=None, initializer=torch.nn.init.xavier_uniform_) self._value_branch = SlimFC( in_size=self.attention_dim, out_size=1, activation_fn=None, initializer=torch.nn.init.xavier_uniform_) self.view_requirements = self.gtrxl.view_requirements
def __init__(self): nn.Module.__init__(self) self.a2_hidden = SlimFC(in_size=1, out_size=16, activation_fn=nn.Tanh, initializer=normc_init_torch(1.0)) self.a2_logits = SlimFC(in_size=16, out_size=2, activation_fn=None, initializer=normc_init_torch(0.01))
def __init__(self, num_states, num_actions, dqn_config): """ Initialize a deep Q-learning network for testing algorithm in_features: number of features of input. num_actions: number of action-value to output, one-to-one correspondence to action in game. """ super(DQNTransitionModule, self).__init__() self.num_states = num_states self.num_actions = num_actions self.device = torch.device(f"cuda:{dqn_config['cuda_id']}" if torch.cuda.is_available() else "cpu") self.state_emb_model = None layers = [] prev_layer_size = num_states self.state_emb_size = 32 for i, size in enumerate([128, 64, self.state_emb_size]): layers.append( SlimFC( in_size=prev_layer_size, out_size=size, initializer=torch_normc_initializer(1.0), activation_fn=nn.ReLU)) prev_layer_size = size self.state_emb_model = nn.Sequential(*layers).to(self.device) self.action_emb_model = None layers = [] prev_layer_size = num_actions self.action_emb_size = 16 for i, size in enumerate([64, self.action_emb_size]): layers.append( SlimFC( in_size=prev_layer_size, out_size=size, initializer=torch_normc_initializer(1.0), activation_fn=nn.ReLU)) prev_layer_size = size self.action_emb_model = nn.Sequential(*layers).to(self.device) self.transition_emb_model = None layers = [] prev_layer_size = self.state_emb_size + self.action_emb_size self.transition_emb_size = self.state_emb_size for i, size in enumerate([128, self.transition_emb_size]): layers.append( SlimFC( in_size=prev_layer_size, out_size=size, initializer=torch_normc_initializer(1.0), activation_fn=nn.ReLU)) prev_layer_size = size self.transition_emb_model = nn.Sequential(*layers).to(self.device)
def __init__(self, device, num_states=4, num_actions=18): """ Initialize a deep Q-learning network for testing algorithm in_features: number of features of input. num_actions: number of action-value to output, one-to-one correspondence to action in game. """ super(DQNActionModule, self).__init__() self.device = device state_layers = [] action_layers = [] self.num_states = num_states self.num_actions = num_actions # Create layers 0 to second-last. state_prev_layer_size = num_states self.state_hidden_out_size = 32 for size in [256, 128, self.state_hidden_out_size]: state_layers.append( SlimFC(in_size=state_prev_layer_size, out_size=size, initializer=torch_normc_initializer(1.0), activation_fn=nn.ReLU)) state_prev_layer_size = size action_prev_layer_size = num_actions self.action_hidden_out_size = 32 for size in [64, self.action_hidden_out_size]: action_layers.append( SlimFC(in_size=action_prev_layer_size, out_size=size, initializer=torch_normc_initializer(1.0), activation_fn=nn.ReLU)) action_prev_layer_size = size layers = [] prev_layer_size = self.state_hidden_out_size + self.action_hidden_out_size self.hidden_out_size = 32 for size in [64, self.hidden_out_size]: layers.append( SlimFC(in_size=prev_layer_size, out_size=size, initializer=torch_normc_initializer(1.0), activation_fn=nn.ReLU)) prev_layer_size = size self._value_branch = SlimFC(in_size=self.hidden_out_size, out_size=1, initializer=torch_normc_initializer(1.0), activation_fn=None) self._state_hidden_layers = nn.Sequential(*state_layers) self._action_hidden_layers = nn.Sequential(*action_layers) self._hidden_layers = nn.Sequential(*layers)
def __init__(self, in_dim, out_dim, num_heads, head_dim, rel_pos_encoder, input_layernorm=False, output_activation=None, **kwargs): """Initializes a RelativeMultiHeadAttention nn.Module object. Args: in_dim (int): out_dim (int): num_heads (int): The number of attention heads to use. Denoted `H` in [2]. head_dim (int): The dimension of a single(!) attention head Denoted `D` in [2]. rel_pos_encoder (: input_layernorm (bool): Whether to prepend a LayerNorm before everything else. Should be True for building a GTrXL. output_activation (Optional[tf.nn.activation]): Optional tf.nn activation function. Should be relu for GTrXL. **kwargs: """ super().__init__(**kwargs) # No bias or non-linearity. self._num_heads = num_heads self._head_dim = head_dim # 3=Query, key, and value inputs. self._qkv_layer = SlimFC( in_size=in_dim, out_size=3 * num_heads * head_dim, use_bias=False) self._linear_layer = SlimFC( in_size=num_heads * head_dim, out_size=out_dim, use_bias=False, activation_fn=output_activation) self._pos_proj = SlimFC( in_size=in_dim, out_size=num_heads * head_dim, use_bias=False) self._uvar = torch.zeros(num_heads, head_dim) self._vvar = torch.zeros(num_heads, head_dim) nn.init.xavier_uniform_(self._uvar) nn.init.xavier_uniform_(self._vvar) self._rel_pos_encoder = rel_pos_encoder self._input_layernorm = None if input_layernorm: self._input_layernorm = torch.nn.LayerNorm(in_dim)
def __init__(self, obs_space: gym.spaces.Space, action_space: gym.spaces.Space, num_outputs: int, model_config: ModelConfigDict, name: str): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) # Nonlinearity for fully connected net (tanh, relu). Default: "tanh" activation = model_config.get("fcnet_activation") # Number of hidden layers for fully connected net. Default: [256, 256] hiddens = [256, 256] # model_config.get("fcnet_hiddens", []) # Whether to skip the final linear layer used to resize the hidden layer # outputs to size `num_outputs`. If True, then the last hidden layer # should already match num_outputs. # no_final_linear = False self.vf_share_layers = model_config.get("vf_share_layers") self.free_log_std = False self._embedd = nn.Embedding( int(obs_space.high[0][-1]) + 1, CARD_EMBEDD_SIZE) # Player Hot Encoded = 3 * Number of Cards Played per trick = 4 # CARD_EMBEDD_SIZE * Number of Cards Played per trick = 4 self._hidden_layers = self._build_hidden_layers( first_layer_size=FIRST_LAYER_SIZE, hiddens=hiddens, activation=activation) self._value_branch_separate = None self._value_embedding = None if not self.vf_share_layers: # Build a parallel set of hidden layers for the value net. self._value_embedding = nn.Embedding( int(obs_space.high[0][-1]) + 1, CARD_EMBEDD_SIZE) self._value_branch_separate = self._build_hidden_layers( first_layer_size=FIRST_LAYER_SIZE, hiddens=hiddens, activation=activation) self._logits = SlimFC(in_size=hiddens[-1], out_size=num_outputs, initializer=normc_initializer(0.01), activation_fn=None) self._value_branch = SlimFC(in_size=hiddens[-1], out_size=1, initializer=normc_initializer(1.0), activation_fn=None) # Holds the current "base" output (before logits layer). self._features = None # Holds the last input, in case value branch is separate. self._cards_in = None self._players_in = None
def __init__(self, obs_space, action_space, num_outputs, model_config, name): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) self.conv1=SlimConv2d(obs_space.shape[0],32,3,2,1) self.conv2=SlimConv2d(32,32,3,2,1) self.conv3=SlimConv2d(32,32,3,2,1) self.conv4=SlimConv2d(32,32,3,2,1) self.fc1=SlimFC(32*6*6,512) self.fc_out=SlimFC(512,num_outputs)
def __init__(self, obs_space, action_space, num_outputs, model_config, name): nn.Module.__init__(self) super().__init__(obs_space, action_space, None, model_config, name) self.cell_size = model_config["lstm_cell_size"] self.time_major = model_config.get("_time_major", False) self.use_prev_action_reward = model_config[ "lstm_use_prev_action_reward"] self.action_dim = int(np.product(action_space.shape)) # Add prev-action/reward nodes to input to LSTM. if self.use_prev_action_reward: self.num_outputs += 1 + self.action_dim self.lstm = nn.LSTM(self.num_outputs, self.cell_size, batch_first=not self.time_major) self.num_outputs = num_outputs # Postprocess LSTM output with another hidden layer and compute values. self._logits_branch = SlimFC(in_size=self.cell_size, out_size=self.num_outputs, activation_fn=None, initializer=torch.nn.init.xavier_uniform_) self._value_branch = SlimFC(in_size=self.cell_size, out_size=1, activation_fn=None, initializer=torch.nn.init.xavier_uniform_) self.inference_view_requirements.update( dict( **{ SampleBatch.OBS: ViewRequirement(shift=0), SampleBatch.PREV_REWARDS: ViewRequirement(SampleBatch.REWARDS, shift=-1), SampleBatch.PREV_ACTIONS: ViewRequirement(SampleBatch.ACTIONS, space=self.action_space, shift=-1), })) for i in range(2): self.inference_view_requirements["state_in_{}".format(i)] = \ ViewRequirement( "state_out_{}".format(i), shift=-1, space=Box(-1.0, 1.0, shape=(self.cell_size,))) self.inference_view_requirements["state_out_{}".format(i)] = \ ViewRequirement( space=Box(-1.0, 1.0, shape=(self.cell_size,)))
def __init__(self, obs_space: gym.spaces.Space, action_space: gym.spaces.Space, num_outputs: int, model_config: ModelConfigDict, name: str): nn.Module.__init__(self) super().__init__(obs_space, action_space, None, model_config, name) self.cell_size = model_config["lstm_cell_size"] self.time_major = model_config.get("_time_major", False) self.use_prev_action = model_config["lstm_use_prev_action"] self.use_prev_reward = model_config["lstm_use_prev_reward"] if isinstance(action_space, Discrete): self.action_dim = action_space.n elif isinstance(action_space, MultiDiscrete): self.action_dim = np.product(action_space.nvec) elif action_space.shape is not None: self.action_dim = int(np.product(action_space.shape)) else: self.action_dim = int(len(action_space)) # Add prev-action/reward nodes to input to LSTM. if self.use_prev_action: self.num_outputs += self.action_dim if self.use_prev_reward: self.num_outputs += 1 self.lstm = nn.LSTM(self.num_outputs, self.cell_size, batch_first=not self.time_major) self.num_outputs = num_outputs # Postprocess LSTM output with another hidden layer and compute values. self._logits_branch = SlimFC(in_size=self.cell_size, out_size=self.num_outputs, activation_fn=None, initializer=torch.nn.init.xavier_uniform_) self._value_branch = SlimFC(in_size=self.cell_size, out_size=1, activation_fn=None, initializer=torch.nn.init.xavier_uniform_) # Add prev-a/r to this model's view, if required. if model_config["lstm_use_prev_action"]: self.inference_view_requirements[SampleBatch.PREV_ACTIONS] = \ ViewRequirement(SampleBatch.ACTIONS, space=self.action_space, data_rel_pos=-1) if model_config["lstm_use_prev_reward"]: self.inference_view_requirements[SampleBatch.PREV_REWARDS] = \ ViewRequirement(SampleBatch.REWARDS, data_rel_pos=-1)
def __init__(self, obs_space, action_space, num_outputs, model_config, name): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) # Base of the model self.model = TorchFC(obs_space, action_space, num_outputs, model_config, name) # Central VF maps (obs, opp_obs, opp_act) -> vf_pred input_size = 6 + 6 + 2 # obs + opp_obs + opp_act self.central_vf_dense = SlimFC(input_size, 16, activation_fn=nn.Tanh) self.central_vf_out = SlimFC(16, 1)
def __init__(self, observation_space, action_space, num_outputs, model_config, name): TorchModelV2.__init__(self, observation_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) # Non-shared initial layer. self.first_layer = SlimFC(int(np.product(observation_space.shape)), 32, activation_fn=nn.ReLU) # Non-shared final layer. self.last_layer = SlimFC(32, self.num_outputs, activation_fn=nn.ReLU) self.vf = SlimFC(32, 1, activation_fn=None)
def __init__(self, obs_space, action_space, num_outputs, model_config, name): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) self.size_in = 3 self.lstm_state_size = 64 self.grab_location = 1 self.lstm = nn.LSTM(self.size_in, self.lstm_state_size, batch_first=True) self.num_outputs = self.grab_location + 7 ## one grab location + seven (position + translation) # Postprocess LSTM output with another hidden layer and compute values. self.linear = SlimFC(self.lstm_state_size, self.num_outputs, activation_fn="tanh") self.value_branch = SlimFC(self.lstm_state_size, 1, activation_fn=None) self._features = None
def __init__(self, obs_space, action_space, num_outputs, model_config, name): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) activation = get_activation_fn( model_config.get("conv_activation"), framework="torch") filters = model_config.get("conv_filters") if not filters: filters = _get_filter_config(obs_space.shape) # no_final_linear = model_config.get("no_final_linear") # vf_share_layers = model_config.get("vf_share_layers") layers = [] (w, h, in_channels) = obs_space.shape in_size = [w, h] for out_channels, kernel, stride in filters[:-1]: padding, out_size = valid_padding(in_size, kernel, [stride, stride]) layers.append( SlimConv2d( in_channels, out_channels, kernel, stride, padding, activation_fn=activation)) in_channels = out_channels in_size = out_size out_channels, kernel, stride = filters[-1] layers.append( SlimConv2d( in_channels, out_channels, kernel, stride, None, activation_fn=activation)) self._convs = nn.Sequential(*layers) self._logits = SlimFC( out_channels, num_outputs, initializer=nn.init.xavier_uniform_) self._value_branch = SlimFC( out_channels, 1, initializer=normc_initializer()) # Holds the current "base" output (before logits layer). self._features = None
def __init__(self, size_in, size_out, hiddens, activations, init_weights, append_log_std=False, log_std_type='constant', sample_std=1.0): super().__init__() layers = [] prev_layer_size = size_in for i, size_hidden in enumerate(hiddens + [size_out]): layers.append( SlimFC(in_size=prev_layer_size, out_size=size_hidden, initializer=normc_initializer(init_weights[i]), activation_fn=get_activation_fn(activations[i], framework="torch"))) prev_layer_size = size_hidden if append_log_std: layers.append( AppendLogStd(type=log_std_type, init_val=np.log(sample_std), dim=size_out)) self._model = nn.Sequential(*layers)
def create_actor(self): model_config = self.model_config layers = [] activation_stage1 = model_config.get("fcnet_activation_stage1") hiddens_stage1 = model_config.get("fcnet_hiddens_stage1") self.gru_cell_size = model_config.get("gru_cell_size") activation_stage2 = model_config.get("fcnet_activation_stage2") hiddens_stage2 = model_config.get("fcnet_hiddens_stage2") prev_layer_size = self.true_obs_space.shape[1] # obs prev_layer_size += self.nbr_agents # one hot encoding of the agent id for size in hiddens_stage1: layers.append( SlimFC(in_size=prev_layer_size, out_size=size, initializer=normc_initializer(1.0), activation_fn=activation_stage1)) prev_layer_size = size stage1 = nn.Sequential(*layers) gru = nn.GRU(input_size=prev_layer_size, hidden_size=self.gru_cell_size, num_layers=1, batch_first=not self.is_time_major()) prev_layer_size = self.gru_cell_size layers = [] for size in hiddens_stage2: layers.append( SlimFC(in_size=prev_layer_size, out_size=size, initializer=normc_initializer(1.0), activation_fn=activation_stage2)) prev_layer_size = size layers.append( SlimFC(in_size=prev_layer_size, out_size=self.nbr_actions, initializer=normc_initializer(1.0))) stage2 = nn.Sequential(*layers) return stage1, gru, stage2
def build_q_net(name_): act = get_activation_fn(critic_hidden_activation, framework="torch") init = nn.init.xavier_uniform_ # For discrete actions, only obs. q_net = nn.Sequential() ins = embed_dim for i, n in enumerate(critic_hiddens): q_net.add_module( "{}_hidden_{}".format(name_, i), SlimFC(ins, n, initializer=init, activation_fn=act)) ins = n q_net.add_module( "{}_out".format(name_), SlimFC(ins, q_outs, initializer=init, activation_fn=None)) return q_net
def __init__(self, in_dim: int, out_dim: int, num_heads: int, head_dim: int, **kwargs): """ in_dim (int): Dimension of input out_dim (int): Dimension of output num_heads (int): Number of attention heads head_dim (int): Output dimension of each attention head """ super().__init__(**kwargs) # No bias or non-linearity. self._num_heads = num_heads self._head_dim = head_dim self._qkv_layer = SlimFC( in_size=in_dim, out_size=3 * num_heads * head_dim, use_bias=False) self._linear_layer = SlimFC( in_size=num_heads * head_dim, out_size=out_dim, use_bias=False)
def create_fc_net(layer_dims, activation): layers = [] for i in range(len(layer_dims) - 1): layers.append( SlimFC(in_size=layer_dims[i], out_size=layer_dims[i + 1], use_bias=False, activation_fn=activation)) return nn.Sequential(*layers)
def __init__( self, obs_space: gym.spaces.Space, action_space: gym.spaces.Space, num_outputs: int, model_config: ModelConfigDict, name: str, **custom_model_kwargs, ): nn.Module.__init__(self) super(DNCMemory, self).__init__(obs_space, action_space, num_outputs, model_config, name) self.num_outputs = num_outputs self.obs_dim = gym.spaces.utils.flatdim(obs_space) self.act_dim = gym.spaces.utils.flatdim(action_space) self.cfg = dict(self.DEFAULT_CONFIG, **custom_model_kwargs) assert (self.cfg["num_layers"] == 1 ), "num_layers != 1 has not been implemented yet" self.cur_val = None self.preprocessor = torch.nn.Sequential( torch.nn.Linear(self.obs_dim, self.cfg["preprocessor_input_size"]), self.cfg["preprocessor"], torch.nn.Linear(self.cfg["preprocessor_output_size"], self.obs_dim)) self.logit_branch = SlimFC( in_size=self.obs_dim, out_size=self.num_outputs, activation_fn=None, initializer=torch.nn.init.xavier_uniform_, ) self.value_branch = SlimFC( in_size=self.obs_dim, out_size=1, activation_fn=None, initializer=torch.nn.init.xavier_uniform_, ) self.dnc: Union[None, DNC] = None
def __init__(self, obs_space, action_space, num_outputs, model_config, name): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) obs_space_ = obs_space.original_space data, privates = obs_space_.spaces['data'], obs_space_.spaces['privates'] N, T, L = data.shape adjusted_data_shape = (T, N*L) activation = model_config.get("fcnet_activation") hiddens = model_config.get("fcnet_hiddens", [100, 100]) lstm_dim = model_config.get("lstm_cell_size", 128) self.lstm_net = LSTM(input_dim=adjusted_data_shape[-1], hidden_dim=lstm_dim, num_layers=2) prev_layer_size = lstm_dim + int(np.product(privates.shape)) layers = [] for size in hiddens: layers.append( SlimFC( in_size=prev_layer_size, out_size=size, initializer=normc_initializer(1.0), activation_fn=activation)) prev_layer_size = size self._hidden_layers = nn.Sequential(*layers) self._features = None self._policy_net = SlimFC( in_size=prev_layer_size, out_size=num_outputs, initializer=normc_initializer(1.0), activation_fn=activation) self._value_net = SlimFC( in_size=prev_layer_size, out_size=1, initializer=normc_initializer(1.0), activation_fn=activation)
def __init__(self, obs_space, action_space, num_outputs, model_config, name): nn.Module.__init__(self) super().__init__(obs_space, action_space, None, model_config, name) self.cell_size = model_config["lstm_cell_size"] self.lstm = nn.LSTM(self.num_outputs, self.cell_size, batch_first=True) self.num_outputs = num_outputs # Postprocess LSTM output with another hidden layer and compute values. self._logits_branch = SlimFC(in_size=self.cell_size, out_size=self.num_outputs, activation_fn=None, initializer=torch.nn.init.xavier_uniform_) self._value_branch = SlimFC(in_size=self.cell_size, out_size=1, activation_fn=None, initializer=torch.nn.init.xavier_uniform_)
def __init__(self, obs_space, action_space, num_outputs, model_config, name): # Pass num_outputs=None into super constructor (so that no action/ # logits output layer is built). # Alternatively, you can pass in num_outputs=[last layer size of # config[model][fcnet_hiddens]] AND set no_last_linear=True, but # this seems more tedious as you will have to explain users of this # class that num_outputs is NOT the size of your Q-output layer. nn.Module.__init__(self) super(TorchDuelingQModel, self).__init__(obs_space, action_space, None, model_config, name) # Now: self.num_outputs contains the last layer's size, which # we can use to construct the dueling head (see torch: SlimFC # below). # Construct advantage head ... self.A = SlimFC(in_size=self.num_outputs, out_size=num_outputs) # ... and value head. self.V = SlimFC(in_size=self.num_outputs, out_size=1)