def default_twinq_net_fn(env):
"""
Returns a default Twinq network
"""
assert isinstance(env.action_space, spaces.Discrete)
if isinstance(env.observation_space, spaces.Box):
obs_shape = env.observation_space.shape
elif isinstance(env.observation_space, spaces.Tuple):
obs_shape = env.observation_space.spaces[0].shape
else:
raise ValueError("Incompatible observation space: {}".format(
env.observation_space))
# Assume CHW observation space
if len(obs_shape) == 1:
model_config = {
"type": "MultiLayerPerceptron",
"in_size": int(obs_shape[0]) + int(env.action_space.n),
"layer_sizes": [64, 64],
}
else:
raise ValueError("Incompatible observation shape: {}".format(
env.observation_space.shape))
model_config["out_size"] = 1
q1 = model_factory(**model_config)
q2 = model_factory(**model_config)
return (q1, q2)
def __init__(
self,
in_size=None,
out_size=None,
presence_feature_idx=0,
embedding_layer_kwargs=None,
attention_layer_kwargs=None,
output_layer_kwargs=None,
**kwargs
):
super().__init__(**kwargs)
self.out_size = out_size
self.presence_feature_idx = presence_feature_idx
embedding_layer_kwargs = embedding_layer_kwargs or {}
if not embedding_layer_kwargs.get("in_size", None):
embedding_layer_kwargs["in_size"] = in_size
self.ego_embedding = model_factory(**embedding_layer_kwargs)
self.embedding = model_factory(**embedding_layer_kwargs)
attention_layer_kwargs = attention_layer_kwargs or {}
self.attention_layer = EgoAttention(**attention_layer_kwargs)
output_layer_kwargs = output_layer_kwargs or {}
output_layer_kwargs["in_size"] = self.attention_layer.feature_size
output_layer_kwargs["out_size"] = self.out_size
self.output_layer = model_factory(**output_layer_kwargs)
def __init__(self,
activation="RELU",
in_channels=None,
in_height=None,
in_width=None,
head_mlp_kwargs=None,
out_size=None,
is_policy=False,
transpose_obs=False,
**kwargs):
super().__init__()
self.activation = activation_factory(activation)
self.conv1 = nn.Conv2d(in_channels, 16, kernel_size=2, stride=2)
self.conv2 = nn.Conv2d(16, 32, kernel_size=2, stride=2)
self.conv3 = nn.Conv2d(32, 64, kernel_size=2, stride=2)
# MLP Head
# Number of Linear input connections depends on output of conv2d layers
# and therefore the input image size, so compute it.
def conv2d_size_out(size, kernel_size=2, stride=2):
return (size - (kernel_size - 1) - 1) // stride + 1
convw = conv2d_size_out(conv2d_size_out(conv2d_size_out(in_width)))
convh = conv2d_size_out(conv2d_size_out(conv2d_size_out(in_height)))
assert convh > 0 and convw > 0
self.head_mlp_kwargs = head_mlp_kwargs or {}
self.head_mlp_kwargs["in_size"] = convw * convh * 64
self.head_mlp_kwargs["out_size"] = out_size
self.head_mlp_kwargs["is_policy"] = is_policy
self.head = model_factory(**self.head_mlp_kwargs)
self.is_policy = is_policy
self.transpose_obs = transpose_obs
def default_policy_net_fn(env):
"""
Returns a default value network.
"""
if isinstance(env.observation_space, spaces.Box):
obs_shape = env.observation_space.shape
elif isinstance(env.observation_space, spaces.Tuple):
obs_shape = env.observation_space.spaces[0].shape
else:
raise ValueError("Incompatible observation space: {}".format(
env.observation_space))
if len(obs_shape) == 3:
if obs_shape[0] < obs_shape[1] and obs_shape[0] < obs_shape[1]:
# Assume CHW observation space
model_config = {
"type": "ConvolutionalNetwork",
"is_policy": True,
"in_channels": int(obs_shape[0]),
"in_height": int(obs_shape[1]),
"in_width": int(obs_shape[2]),
}
elif obs_shape[2] < obs_shape[0] and obs_shape[2] < obs_shape[1]:
# Assume WHC observation space
model_config = {
"type": "ConvolutionalNetwork",
"is_policy": True,
"transpose_obs": True,
"in_channels": int(obs_shape[2]),
"in_height": int(obs_shape[1]),
"in_width": int(obs_shape[0]),
}
elif len(obs_shape) == 2:
model_config = {
"type": "ConvolutionalNetwork",
"is_policy": True,
"in_channels": int(1),
"in_height": int(obs_shape[0]),
"in_width": int(obs_shape[1]),
}
elif len(obs_shape) == 1:
model_config = {
"type": "MultiLayerPerceptron",
"in_size": int(obs_shape[0]),
"layer_sizes": [64, 64],
"reshape": False,
"is_policy": True,
}
else:
raise ValueError("Incompatible observation shape: {}".format(
env.observation_space.shape))
if isinstance(env.action_space, spaces.Discrete):
model_config["out_size"] = env.action_space.n
elif isinstance(env.action_space, spaces.Tuple):
model_config["out_size"] = env.action_space.spaces[0].n
return model_factory(**model_config)
def default_q_net_fn(env, **kwargs):
"""
Returns a default Q value network.
"""
del kwargs
model_config = {
"type": "MultiLayerPerceptron",
"layer_sizes": (64, 64),
"reshape": False,
}
model_config = size_model_config(env, **model_config)
return model_factory(**model_config)
def __init__(
self,
in_size=None,
base_module_kwargs=None,
value_kwargs=None,
advantage_kwargs=None,
out_size=None,
):
super().__init__()
self.out_size = out_size
base_module_kwargs = base_module_kwargs or {}
base_module_kwargs["in_size"] = in_size
self.base_module = model_factory(**base_module_kwargs)
value_kwargs = value_kwargs or {}
value_kwargs["in_size"] = self.base_module.layer_sizes[-1]
value_kwargs["out_size"] = 1
self.value = model_factory(**value_kwargs)
advantage_kwargs = advantage_kwargs or {}
advantage_kwargs["in_size"] = self.base_module.layer_sizes[-1]
advantage_kwargs["out_size"] = out_size
self.advantage = model_factory(**advantage_kwargs)
def default_value_net_fn(env):
"""
Returns a default value network.
"""
if isinstance(env.observation_space, spaces.Box):
obs_shape = env.observation_space.shape
elif isinstance(env.observation_space, spaces.Tuple):
obs_shape = env.observation_space.spaces[0].shape
else:
raise ValueError("Incompatible observation space: {}".format(
env.observation_space))
# Assume CHW observation space
if len(obs_shape) == 3:
model_config = {
"type": "ConvolutionalNetwork",
"in_channels": int(obs_shape[0]),
"in_height": int(obs_shape[1]),
"in_width": int(obs_shape[2]),
}
elif len(obs_shape) == 2:
model_config = {
"type": "ConvolutionalNetwork",
"in_channels": int(1),
"in_height": int(obs_shape[0]),
"in_width": int(obs_shape[1]),
}
elif len(obs_shape) == 1:
model_config = {
"type": "MultiLayerPerceptron",
"in_size": int(obs_shape[0]),
"layer_sizes": [64, 64],
}
else:
raise ValueError("Incompatible observation shape: {}".format(
env.observation_space.shape))
model_config["out_size"] = 1
return model_factory(**model_config)