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):
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 conv_layers(x, model_config, obs_space, prefix=""):
filters = model_config.get("conv_filters")
if not filters:
filters = _get_filter_config(obs_space.shape)
activation = get_conv_activation(model_config)
for i, (out_size, kernel, stride) in enumerate(filters, 1):
x = tf.keras.layers.Conv2D(
out_size,
kernel,
strides=(stride, stride),
activation=activation,
padding="same",
data_format="channels_last",
name=f"{prefix}conv{i}",
)(x)
return x
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
if not model_config.get("conv_filters"):
model_config["conv_filters"] = _get_filter_config(obs_space.shape)
TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
model_config, name)
nn.Module.__init__(self)
activation = self.model_config.get("conv_activation")
filters = self.model_config["conv_filters"]
no_final_linear = self.model_config.get("no_final_linear")
vf_share_layers = self.model_config.get("vf_share_layers")
# Whether the last layer is the output of a Flattened (rather than
# a n x (1,1) Conv2D).
self.last_layer_is_flattened = False
self._logits = None
layers = []
(w, h, in_channels) = obs_space.shape
in_size = [w, h]
for out_channels, kernel, stride in filters[:-1]:
padding, out_size = same_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]
# No final linear: Last layer is a Conv2D and uses num_outputs.
if no_final_linear and num_outputs:
layers.append(
SlimConv2d(
in_channels,
num_outputs,
kernel,
stride,
None, # padding=valid
activation_fn=activation))
out_channels = num_outputs
# Finish network normally (w/o overriding last layer size with
# `num_outputs`), then add another linear one of size `num_outputs`.
else:
layers.append(
SlimConv2d(
in_channels,
out_channels,
kernel,
stride,
None, # padding=valid
activation_fn=activation))
# num_outputs defined. Use that to create an exact
# `num_output`-sized (1,1)-Conv2D.
if num_outputs:
in_size = [
np.ceil((in_size[0] - kernel[0]) / stride),
np.ceil((in_size[1] - kernel[1]) / stride)
]
padding, _ = same_padding(in_size, [1, 1], [1, 1])
self._logits = SlimConv2d(
out_channels,
num_outputs, [1, 1],
1,
padding,
activation_fn=None)
# num_outputs not known -> Flatten, then set self.num_outputs
# to the resulting number of nodes.
else:
self.last_layer_is_flattened = True
layers.append(nn.Flatten())
self.num_outputs = out_channels
self._convs = nn.Sequential(*layers)
# Build the value layers
self._value_branch_separate = self._value_branch = None
if vf_share_layers:
self._value_branch = SlimFC(
out_channels,
1,
initializer=normc_initializer(0.01),
activation_fn=None)
else:
vf_layers = []
(w, h, in_channels) = obs_space.shape
in_size = [w, h]
for out_channels, kernel, stride in filters[:-1]:
padding, out_size = same_padding(in_size, kernel,
[stride, stride])
vf_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]
vf_layers.append(
SlimConv2d(
in_channels,
out_channels,
kernel,
stride,
None,
activation_fn=activation))
vf_layers.append(
SlimConv2d(
in_channels=out_channels,
out_channels=1,
kernel=1,
stride=1,
padding=None,
activation_fn=None))
self._value_branch_separate = nn.Sequential(*vf_layers)
# Holds the current "base" output (before logits layer).
self._features = None
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
super(VisionNetwork, self).__init__(obs_space, action_space,
num_outputs, model_config, name)
activation = get_activation_fn(model_config.get("conv_activation"))
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")
inputs = tf.keras.layers.Input(shape=obs_space.shape,
name="observations")
last_layer = inputs
# Build the action layers
for i, (out_size, kernel, stride) in enumerate(filters[:-1], 1):
last_layer = tf.keras.layers.Conv2D(
out_size,
kernel,
strides=(stride, stride),
activation=activation,
padding="same",
name="conv{}".format(i))(last_layer)
out_size, kernel, stride = filters[-1]
if no_final_linear:
# the last layer is adjusted to be of size num_outputs
last_layer = tf.keras.layers.Conv2D(num_outputs,
kernel,
strides=(stride, stride),
activation=activation,
padding="valid",
name="conv_out")(last_layer)
conv_out = last_layer
else:
last_layer = tf.keras.layers.Conv2D(
out_size,
kernel,
strides=(stride, stride),
activation=activation,
padding="valid",
name="conv{}".format(i + 1))(last_layer)
conv_out = tf.keras.layers.Conv2D(num_outputs, [1, 1],
activation=None,
padding="same",
name="conv_out")(last_layer)
# Build the value layers
if vf_share_layers:
last_layer = tf.keras.layers.Lambda(
lambda x: tf.squeeze(x, axis=[1, 2]))(last_layer)
value_out = tf.keras.layers.Dense(
1,
name="value_out",
activation=None,
kernel_initializer=normc_initializer(0.01))(last_layer)
else:
# build a parallel set of hidden layers for the value net
last_layer = inputs
for i, (out_size, kernel, stride) in enumerate(filters[:-1], 1):
last_layer = tf.keras.layers.Conv2D(
out_size,
kernel,
strides=(stride, stride),
activation=activation,
padding="same",
name="conv_value_{}".format(i))(last_layer)
out_size, kernel, stride = filters[-1]
last_layer = tf.keras.layers.Conv2D(
out_size,
kernel,
strides=(stride, stride),
activation=activation,
padding="valid",
name="conv_value_{}".format(i + 1))(last_layer)
last_layer = tf.keras.layers.Conv2D(
1, [1, 1],
activation=None,
padding="same",
name="conv_value_out")(last_layer)
value_out = tf.keras.layers.Lambda(
lambda x: tf.squeeze(x, axis=[1, 2]))(last_layer)
self.base_model = tf.keras.Model(inputs, [conv_out, value_out])
self.register_variables(self.base_model.variables)
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
super(VisionNetwork, self).__init__(obs_space, action_space,
num_outputs, model_config, name)
activation = get_activation_fn(model_config.get("conv_activation"))
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")
inputs = tf.keras.layers.Input(shape=obs_space.shape,
name="observations")
last_layer = inputs
# Whether the last layer is the output of a Flattened (rather than
# a n x (1,1) Conv2D).
self.last_layer_is_flattened = False
# Build the action layers
for i, (out_size, kernel, stride) in enumerate(filters[:-1], 1):
last_layer = tf.keras.layers.Conv2D(
out_size,
kernel,
strides=(stride, stride),
activation=activation,
padding="same",
data_format="channels_last",
name="conv{}".format(i))(last_layer)
out_size, kernel, stride = filters[-1]
# No final linear: Last layer is a Conv2D and uses num_outputs.
if no_final_linear and num_outputs:
last_layer = tf.keras.layers.Conv2D(num_outputs,
kernel,
strides=(stride, stride),
activation=activation,
padding="valid",
data_format="channels_last",
name="conv_out")(last_layer)
conv_out = last_layer
# Finish network normally (w/o overriding last layer size with
# `num_outputs`), then add another linear one of size `num_outputs`.
else:
last_layer = tf.keras.layers.Conv2D(
out_size,
kernel,
strides=(stride, stride),
activation=activation,
padding="valid",
data_format="channels_last",
name="conv{}".format(i + 1))(last_layer)
# num_outputs defined. Use that to create an exact
# `num_output`-sized (1,1)-Conv2D.
if num_outputs:
conv_out = tf.keras.layers.Conv2D(num_outputs, [1, 1],
activation=None,
padding="same",
data_format="channels_last",
name="conv_out")(last_layer)
# num_outputs not known -> Flatten, then set self.num_outputs
# to the resulting number of nodes.
else:
self.last_layer_is_flattened = True
conv_out = tf.keras.layers.Flatten(
data_format="channels_last")(last_layer)
self.num_outputs = conv_out.shape[1]
# Build the value layers
if vf_share_layers:
last_layer = tf.keras.layers.Lambda(
lambda x: tf.squeeze(x, axis=[1, 2]))(last_layer)
value_out = tf.keras.layers.Dense(
1,
name="value_out",
activation=None,
kernel_initializer=normc_initializer(0.01))(last_layer)
else:
# build a parallel set of hidden layers for the value net
last_layer = inputs
for i, (out_size, kernel, stride) in enumerate(filters[:-1], 1):
last_layer = tf.keras.layers.Conv2D(
out_size,
kernel,
strides=(stride, stride),
activation=activation,
padding="same",
data_format="channels_last",
name="conv_value_{}".format(i))(last_layer)
out_size, kernel, stride = filters[-1]
last_layer = tf.keras.layers.Conv2D(
out_size,
kernel,
strides=(stride, stride),
activation=activation,
padding="valid",
data_format="channels_last",
name="conv_value_{}".format(i + 1))(last_layer)
last_layer = tf.keras.layers.Conv2D(
1, [1, 1],
activation=None,
padding="same",
data_format="channels_last",
name="conv_value_out")(last_layer)
value_out = tf.keras.layers.Lambda(
lambda x: tf.squeeze(x, axis=[1, 2]))(last_layer)
self.base_model = tf.keras.Model(inputs, [conv_out, value_out])
self.register_variables(self.base_model.variables)
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
with tf.variable_scope(f"{name}_model", reuse=tf.AUTO_REUSE):
super().__init__(obs_space, action_space, num_outputs,
model_config, name)
custom_opts = model_config.get("custom_options", {})
self.use_comm = custom_opts.get("use_comm", True)
self.message_coeff = custom_opts.get("message_entropy_coeff", 0.0)
obs_space_shape = custom_opts.get("obs_shape", obs_space.shape)
if len(obs_space.shape) == 1:
n_extra_obs = obs_space.shape[0] - np.prod(obs_space_shape)
else:
n_extra_obs = 0
# Conv
activation = get_activation_fn(model_config.get("conv_activation"))
filters = model_config.get("conv_filters")
if filters is None:
filters = _get_filter_config(obs_space_shape)
inputs = tf.keras.layers.Input(shape=(None, *obs_space_shape),
name=f"{name}_observations_time")
model_inputs = [inputs]
cnn_in = tf.reshape(inputs, [-1, *obs_space_shape])
conv_out = build_cnn(cnn_in,
filters,
activation,
name=f"{name}_conv")
# FC
activation = get_activation_fn(
model_config.get("fcnet_activation"))
hiddens = model_config.get("fcnet_hiddens")
if n_extra_obs > 0:
extra_inputs = tf.keras.layers.Input(
shape=(n_extra_obs, ), name=f"{name}_extra_observations")
model_inputs.append(extra_inputs)
fc_in = tf.keras.layers.Concatenate(name=f"{name}_fc_in")(
[tf.keras.layers.Flatten()(conv_out), extra_inputs])
else:
fc_in = tf.keras.layers.Flatten(name=f"{name}_fc_in")(conv_out)
fc_out = build_fc(fc_in, hiddens, activation, name=f"{name}_fc")
# LSTM
self.cell_size = model_config.get("lstm_cell_size", 256)
state_in_h = tf.keras.layers.Input(shape=(self.cell_size, ),
name=f"{name}_h")
state_in_c = tf.keras.layers.Input(shape=(self.cell_size, ),
name=f"{name}_c")
seq_in = tf.keras.layers.Input(shape=(),
name=f"{name}_seq_in",
dtype=tf.int32)
prev_actions = tf.keras.layers.Input(shape=(),
name=f"{name}_prev_actions",
dtype=tf.int32)
prev_rewards = tf.keras.layers.Input(shape=(),
name=f"{name}_prev_rewards")
model_inputs.extend(
[prev_actions, prev_rewards, seq_in, state_in_h, state_in_c])
if model_config.get("lstm_use_prev_action_reward"):
prev_actions_onehot = tf.one_hot(prev_actions,
action_space[0].n)
in_tensors = [fc_out, prev_actions_onehot, prev_rewards]
else:
in_tensors = [fc_out]
# CPC objective
self.use_cpc = custom_opts.get("use_cpc", False)
if self.use_cpc:
cpc_params = custom_opts["cpc_opts"]
self.cpc_in_shape = [cpc_params["cpc_code_size"]]
self.cpc_out_shape = [
cpc_params["cpc_len"], cpc_params["cpc_code_size"]
]
cpc_params["name"] = f"{name}_cpc"
# The actual CPC encodings
self._cpc_ins = None
self._cpc_preds = None
else:
cpc_params = {}
lstm_out, model_outputs = build_lstm(
in_tensors,
state_in_h=state_in_h,
state_in_c=state_in_c,
seq_in=seq_in,
cell_size=self.cell_size,
add_cpc=self.use_cpc,
cpc_params=cpc_params,
name=f"{name}_lstm",
)
# Final layer, logits has both actions and messages
self.use_inference_policy = custom_opts.get(
"use_inference_policy", False)
if self.use_inference_policy:
inference_policy_opts = custom_opts["inference_policy_opts"]
self.pm_type = inference_policy_opts["type"]
self.ewma_momentum = inference_policy_opts.get("ewma_momentum")
self.pm_hidden = inference_policy_opts.get(
"pm_hidden", [64, 64])
self.message_size = action_space[1].n
action_logits = tf.keras.layers.Dense(
action_space[0].n,
activation=tf.keras.activations.linear,
name=f"{name}_action_logits",
)(lstm_out)
unscaled_message_logits = tf.keras.layers.Dense(
self.message_size,
activation=tf.keras.activations.linear,
name=f"{name}_unscaled_message_logits",
)(lstm_out)
unscaled_message_p = tf.nn.softmax(unscaled_message_logits)
model_outputs.append(unscaled_message_p)
if self.pm_type == "moving_avg":
self._avg_message_p = tf.Variable(
name=f"{name}_avg",
initial_value=tf.ones(
(self.message_size, )) / self.message_size,
trainable=False,
)
avg_message_vars = [self._avg_message_p]
if self.ewma_momentum is None:
self._avg_message_t = tf.Variable(
name=f"{name}_t",
initial_value=tf.zeros(()),
trainable=False,
)
avg_message_vars.append(self._avg_message_t)
self.register_variables(avg_message_vars)
logits = tf.keras.layers.Concatenate(name=f"{name}_logits")(
[action_logits, unscaled_message_logits])
else:
logits = tf.keras.layers.Dense(
num_outputs,
activation=tf.keras.activations.linear,
name=f"{name}_logits",
)(lstm_out)
values = tf.keras.layers.Dense(1,
activation=None,
name=f"{name}_values")(lstm_out)
self._value_out = None # The actual value
model_outputs = [logits, values] + model_outputs
# Create the RNN model
self.rnn_model = tf.keras.Model(inputs=model_inputs,
outputs=model_outputs)
self.register_variables(self.rnn_model.variables)
self._model_out = None # Actual logits
self.rnn_model.summary()
if self.use_inference_policy and self.pm_type == "hyper_nn":
flattened_vars = []
message_model = tf.keras.Model(inputs=model_inputs,
outputs=unscaled_message_logits)
for e in message_model.variables:
flattened_vars.append(
tf.reshape(tf.stop_gradient(e), shape=[1, -1]))
concat_vars = tf.keras.layers.Concatenate()(flattened_vars)
pm_fc_out, pm_fc_vars = build_fc(concat_vars,
self.pm_hidden,
"relu",
name="pm_fc",
return_vars=True)
pm_logits_layer = tf.keras.layers.Dense(
self.message_size,
activation=tf.keras.activations.linear,
name=f"{name}_pm_logits",
)
self._pm_logits = pm_logits_layer(pm_fc_out)
self.register_variables(pm_fc_vars)
self.register_variables(pm_logits_layer.variables)
# Extra variable definitions
self.use_receiver_bias = custom_opts.get("use_receiver_bias",
False)
self.no_message_outputs = None
self._unscaled_message_p = None