def make_networks(
action_spec: specs.BoundedArray,
policy_layer_sizes: Sequence[int] = (256, 256, 256),
critic_layer_sizes: Sequence[int] = (512, 512, 256),
) -> Dict[str, types.TensorTransformation]:
"""Creates networks used by the agent."""
num_dimensions = np.prod(action_spec.shape, dtype=int)
policy_network = snt.Sequential([
networks.LayerNormMLP(policy_layer_sizes, activate_final=True),
networks.NearZeroInitializedLinear(num_dimensions),
networks.TanhToSpec(action_spec)
])
critic_network = snt.Sequential([
# The multiplexer concatenates the observations/actions.
networks.CriticMultiplexer(),
networks.LayerNormMLP(critic_layer_sizes, activate_final=True),
networks.NearZeroInitializedLinear(1),
])
return {
'policy': policy_network,
'critic': critic_network,
'observation': tf2_utils.batch_concat,
}
def make_d4pg_networks(
action_spec,
policy_layer_sizes=(256, 256, 256),
critic_layer_sizes=(512, 512, 256),
vmin=-150.,
vmax=150.,
num_atoms=201):
"""Creates networks used by the d4pg agent."""
num_dimensions = np.prod(action_spec.shape, dtype=int)
policy_layer_sizes = list(policy_layer_sizes) + [int(num_dimensions)]
policy_network = snt.Sequential([
networks.LayerNormMLP(policy_layer_sizes),
networks.TanhToSpec(action_spec)
])
# The multiplexer concatenates the (maybe transformed) observations/actions.
critic_network = snt.Sequential([
networks.CriticMultiplexer(
critic_network=networks.LayerNormMLP(
critic_layer_sizes, activate_final=True)),
networks.DiscreteValuedHead(vmin, vmax, num_atoms)
])
return {
'policy': policy_network,
'critic': critic_network,
'observation': tf2_utils.batch_concat,
}
def make_networks(action_spec: specs.BoundedArray):
"""Creates simple networks for testing.."""
num_dimensions = np.prod(action_spec.shape, dtype=int)
# Create the observation network shared between the policy and critic.
observation_network = tf2_utils.batch_concat
# Create the policy network (head) and the evaluation network.
policy_network = snt.Sequential([
networks.LayerNormMLP([50], activate_final=True),
networks.NearZeroInitializedLinear(num_dimensions),
networks.TanhToSpec(action_spec)
])
evaluator_network = snt.Sequential([observation_network, policy_network])
# Create the critic network.
critic_network = snt.Sequential([
# The multiplexer concatenates the observations/actions.
networks.CriticMultiplexer(),
networks.LayerNormMLP([50], activate_final=True),
networks.NearZeroInitializedLinear(1),
])
return {
'policy': policy_network,
'critic': critic_network,
'observation': observation_network,
'evaluator': evaluator_network,
}
def make_networks(
action_spec: specs.BoundedArray,
policy_layer_sizes: Sequence[int] = (256, 256, 256),
critic_layer_sizes: Sequence[int] = (512, 512, 256),
vmin: float = -150.,
vmax: float = 150.,
num_atoms: int = 51,
) -> Dict[str, Union[snt.Module, Callable[[tf.Tensor], tf.Tensor]]]:
"""Creates networks used by the agent."""
num_dimensions = np.prod(action_spec.shape, dtype=int)
policy_network = snt.Sequential([
networks.LayerNormMLP(policy_layer_sizes, activate_final=True),
networks.NearZeroInitializedLinear(num_dimensions),
networks.TanhToSpec(action_spec)
])
# The multiplexer concatenates the (maybe transformed) observations/actions.
critic_network = snt.Sequential([
networks.CriticMultiplexer(),
networks.LayerNormMLP(critic_layer_sizes, activate_final=True),
networks.DiscreteValuedHead(vmin, vmax, num_atoms),
])
return {
'policy': policy_network,
'critic': critic_network,
'observation': tf2_utils.batch_concat,
}
def __init__(self, visualRadius, action_size, action_spec):
super(PolicyNetwork, self).__init__(name="commons-policy")
self.policy_modules = [
tf2_utils.batch_concat,
snt.Conv2D(3, 4, 4, 2, 1),
tf.nn.relu(),
snt.Conv2D(3, 8, 4, 2, 1),
tf.nn.relu(),
tf.keras.layers.Flatten(),
snt.Linear(8*((ceil(visualRadius/4))**2), 128),
tf.nn.relu(),
snt.Linear(128, action_size),
networks.TanhToSpec(spec=action_spec)
]
self.model = snt.Sequential(self.policy_modules)
def make_d4pg_agent(env_spec: specs.EnvironmentSpec, logger: Logger, checkpoint_path: str, hyperparams: Dict):
params = DEFAULT_PARAMS.copy()
params.update(hyperparams)
action_size = np.prod(env_spec.actions.shape, dtype=int).item()
policy_network = snt.Sequential([
networks.LayerNormMLP(layer_sizes=[*params.pop('policy_layers'), action_size]),
networks.NearZeroInitializedLinear(output_size=action_size),
networks.TanhToSpec(env_spec.actions),
])
critic_network = snt.Sequential([
networks.CriticMultiplexer(
critic_network=networks.LayerNormMLP(layer_sizes=[*params.pop('critic_layers'), 1])
),
networks.DiscreteValuedHead(vmin=-100.0, vmax=100.0, num_atoms=params.pop('atoms'))
])
observation_network = tf.identity
# Make sure observation network is a Sonnet Module.
observation_network = tf2_utils.to_sonnet_module(observation_network)
actor = FeedForwardActor(policy_network=snt.Sequential([
observation_network,
policy_network
]))
# Create optimizers.
policy_optimizer = Adam(params.pop('policy_lr'))
critic_optimizer = Adam(params.pop('critic_lr'))
# The learner updates the parameters (and initializes them).
agent = D4PG(
environment_spec=env_spec,
policy_network=policy_network,
critic_network=critic_network,
observation_network=observation_network,
policy_optimizer=policy_optimizer,
critic_optimizer=critic_optimizer,
logger=logger,
checkpoint_path=checkpoint_path,
**params
)
agent.__setattr__('eval_actor', actor)
return agent
def make_networks(action_spec: specs.BoundedArray):
"""Simple networks for testing.."""
num_dimensions = np.prod(action_spec.shape, dtype=int)
policy_network = snt.Sequential([
networks.LayerNormMLP([50], activate_final=True),
networks.NearZeroInitializedLinear(num_dimensions),
networks.TanhToSpec(action_spec)
])
# The multiplexer concatenates the (maybe transformed) observations/actions.
critic_network = snt.Sequential([
networks.CriticMultiplexer(
critic_network=networks.LayerNormMLP([50], activate_final=True)),
networks.DiscreteValuedHead(-1., 1., 10)
])
return {
'policy': policy_network,
'critic': critic_network,
'observation': tf2_utils.batch_concat,
}
def make_networks(
action_spec: specs.BoundedArray,
policy_layer_sizes: Sequence[int] = (256, 256, 256),
critic_layer_sizes: Sequence[int] = (512, 512, 256),
vmin: float = -150.,
vmax: float = 150.,
num_atoms: int = 51,
) -> Mapping[str, types.TensorTransformation]:
"""Creates networks used by the agent."""
# Get total number of action dimensions from action spec.
num_dimensions = np.prod(action_spec.shape, dtype=int)
# Create the shared observation network; here simply a state-less operation.
observation_network = tf2_utils.batch_concat
# Create the policy network.
policy_network = snt.Sequential([
networks.LayerNormMLP(policy_layer_sizes, activate_final=True),
networks.NearZeroInitializedLinear(num_dimensions),
networks.TanhToSpec(action_spec),
])
# Create the critic network.
critic_network = snt.Sequential([
# The multiplexer concatenates the observations/actions.
networks.CriticMultiplexer(),
networks.LayerNormMLP(critic_layer_sizes, activate_final=True),
networks.DiscreteValuedHead(vmin, vmax, num_atoms),
])
return {
'policy': policy_network,
'critic': critic_network,
'observation': observation_network,
}
environment_spec = specs.make_environment_spec(environment)
#@title Build agent networks
# BUILDING A D4PG AGENT
# Get total number of action dimensions from action spec.
num_dimensions = np.prod(environment_spec.actions.shape, dtype=int)
# Create the shared observation network; here simply a state-less operation.
observation_network = tf2_utils.batch_concat
# Create the deterministic policy network.
policy_network = snt.Sequential([
networks.LayerNormMLP((256, 256, 256), activate_final=True),
networks.NearZeroInitializedLinear(num_dimensions),
networks.TanhToSpec(environment_spec.actions),
])
# Create the distributional critic network.
critic_network = snt.Sequential([
# The multiplexer concatenates the observations/actions.
networks.CriticMultiplexer(),
networks.LayerNormMLP((512, 512, 256), activate_final=True),
networks.DiscreteValuedHead(vmin=-150., vmax=150., num_atoms=51),
])
# Create a logger for the agent and environment loop.
agent_logger = loggers.TerminalLogger(label='agent', time_delta=10.)
env_loop_logger = loggers.TerminalLogger(label='env_loop', time_delta=10.)
# Create the D4PG agent.
def make_networks(
environment_spec: mava_specs.MAEnvironmentSpec,
policy_networks_layer_sizes: Union[Dict[str, Sequence], Sequence] = (
256,
256,
256,
),
critic_networks_layer_sizes: Union[Dict[str, Sequence], Sequence] = (512, 512, 256),
shared_weights: bool = True,
sigma: float = 0.3,
) -> Mapping[str, types.TensorTransformation]:
"""Creates networks used by the agents."""
specs = environment_spec.get_agent_specs()
# Create agent_type specs
if shared_weights:
type_specs = {key.split("_")[0]: specs[key] for key in specs.keys()}
specs = type_specs
if isinstance(policy_networks_layer_sizes, Sequence):
policy_networks_layer_sizes = {
key: policy_networks_layer_sizes for key in specs.keys()
}
if isinstance(critic_networks_layer_sizes, Sequence):
critic_networks_layer_sizes = {
key: critic_networks_layer_sizes for key in specs.keys()
}
observation_networks = {}
policy_networks = {}
critic_networks = {}
for key in specs.keys():
# Get total number of action dimensions from action spec.
num_dimensions = np.prod(specs[key].actions.shape, dtype=int)
# Create the shared observation network; here simply a state-less operation.
observation_network = tf2_utils.to_sonnet_module(tf.identity)
# Create the policy network.
policy_network = snt.Sequential(
[
networks.LayerNormMLP(
policy_networks_layer_sizes[key], activate_final=True
),
networks.NearZeroInitializedLinear(num_dimensions),
networks.TanhToSpec(specs[key].actions),
networks.ClippedGaussian(sigma),
networks.ClipToSpec(specs[key].actions),
]
)
# Create the critic network.
critic_network = snt.Sequential(
[
# The multiplexer concatenates the observations/actions.
networks.CriticMultiplexer(),
networks.LayerNormMLP(
critic_networks_layer_sizes[key], activate_final=False
),
snt.Linear(1),
]
)
observation_networks[key] = observation_network
policy_networks[key] = policy_network
critic_networks[key] = critic_network
return {
"policies": policy_networks,
"critics": critic_networks,
"observations": observation_networks,
}
def make_networks(
environment_spec: mava_specs.MAEnvironmentSpec,
policy_networks_layer_sizes: Union[Dict[str, Sequence], Sequence] = (
256,
256,
256,
),
critic_networks_layer_sizes: Union[Dict[str, Sequence],
Sequence] = (512, 512, 256),
shared_weights: bool = True,
) -> Dict[str, snt.Module]:
"""Creates networks used by the agents."""
# Create agent_type specs.
specs = environment_spec.get_agent_specs()
if shared_weights:
type_specs = {key.split("_")[0]: specs[key] for key in specs.keys()}
specs = type_specs
if isinstance(policy_networks_layer_sizes, Sequence):
policy_networks_layer_sizes = {
key: policy_networks_layer_sizes
for key in specs.keys()
}
if isinstance(critic_networks_layer_sizes, Sequence):
critic_networks_layer_sizes = {
key: critic_networks_layer_sizes
for key in specs.keys()
}
observation_networks = {}
policy_networks = {}
critic_networks = {}
for key in specs.keys():
# Create the shared observation network; here simply a state-less operation.
observation_network = tf2_utils.to_sonnet_module(tf.identity)
# Note: The discrete case must be placed first as it inherits from BoundedArray.
if isinstance(specs[key].actions,
dm_env.specs.DiscreteArray): # discrete
num_actions = specs[key].actions.num_values
policy_network = snt.Sequential([
networks.LayerNormMLP(
tuple(policy_networks_layer_sizes[key]) + (num_actions, ),
activate_final=False,
),
tf.keras.layers.Lambda(lambda logits: tfp.distributions.
Categorical(logits=logits)),
])
elif isinstance(specs[key].actions,
dm_env.specs.BoundedArray): # continuous
num_actions = np.prod(specs[key].actions.shape, dtype=int)
policy_network = snt.Sequential([
networks.LayerNormMLP(policy_networks_layer_sizes[key],
activate_final=True),
networks.MultivariateNormalDiagHead(
num_dimensions=num_actions),
networks.TanhToSpec(specs[key].actions),
])
else:
raise ValueError(
f"Unknown action_spec type, got {specs[key].actions}.")
critic_network = snt.Sequential([
networks.LayerNormMLP(critic_networks_layer_sizes[key],
activate_final=True),
networks.NearZeroInitializedLinear(1),
])
observation_networks[key] = observation_network
policy_networks[key] = policy_network
critic_networks[key] = critic_network
return {
"policies": policy_networks,
"critics": critic_networks,
"observations": observation_networks,
}
def make_default_networks(
environment_spec: mava_specs.MAEnvironmentSpec,
policy_networks_layer_sizes: Union[Dict[str, Sequence], Sequence] = (256, 256, 256),
critic_networks_layer_sizes: Union[Dict[str, Sequence], Sequence] = (512, 512, 256),
shared_weights: bool = True,
) -> Dict[str, snt.Module]:
"""Default networks for mappo.
Args:
environment_spec (mava_specs.MAEnvironmentSpec): description of the action and
observation spaces etc. for each agent in the system.
policy_networks_layer_sizes (Union[Dict[str, Sequence], Sequence], optional):
size of policy networks. Defaults to (256, 256, 256).
critic_networks_layer_sizes (Union[Dict[str, Sequence], Sequence], optional):
size of critic networks. Defaults to (512, 512, 256).
shared_weights (bool, optional): whether agents should share weights or not.
Defaults to True.
Raises:
ValueError: Unknown action_spec type, if actions aren't DiscreteArray
or BoundedArray.
Returns:
Dict[str, snt.Module]: returned agent networks.
"""
# Create agent_type specs.
specs = environment_spec.get_agent_specs()
if shared_weights:
type_specs = {key.split("_")[0]: specs[key] for key in specs.keys()}
specs = type_specs
if isinstance(policy_networks_layer_sizes, Sequence):
policy_networks_layer_sizes = {
key: policy_networks_layer_sizes for key in specs.keys()
}
if isinstance(critic_networks_layer_sizes, Sequence):
critic_networks_layer_sizes = {
key: critic_networks_layer_sizes for key in specs.keys()
}
observation_networks = {}
policy_networks = {}
critic_networks = {}
for key in specs.keys():
# Create the shared observation network; here simply a state-less operation.
observation_network = tf2_utils.to_sonnet_module(tf.identity)
# Note: The discrete case must be placed first as it inherits from BoundedArray.
if isinstance(specs[key].actions, dm_env.specs.DiscreteArray): # discrete
num_actions = specs[key].actions.num_values
policy_network = snt.Sequential(
[
networks.LayerNormMLP(
tuple(policy_networks_layer_sizes[key]) + (num_actions,),
activate_final=False,
),
tf.keras.layers.Lambda(
lambda logits: tfp.distributions.Categorical(logits=logits)
),
]
)
elif isinstance(specs[key].actions, dm_env.specs.BoundedArray): # continuous
num_actions = np.prod(specs[key].actions.shape, dtype=int)
policy_network = snt.Sequential(
[
networks.LayerNormMLP(
policy_networks_layer_sizes[key], activate_final=True
),
networks.MultivariateNormalDiagHead(num_dimensions=num_actions),
networks.TanhToSpec(specs[key].actions),
]
)
else:
raise ValueError(f"Unknown action_spec type, got {specs[key].actions}.")
critic_network = snt.Sequential(
[
networks.LayerNormMLP(
list(critic_networks_layer_sizes[key]) + [1], activate_final=False
),
]
)
observation_networks[key] = observation_network
policy_networks[key] = policy_network
critic_networks[key] = critic_network
return {
"policies": policy_networks,
"critics": critic_networks,
"observations": observation_networks,
}
def make_default_networks(
environment_spec: mava_specs.MAEnvironmentSpec,
policy_networks_layer_sizes: Union[Dict[str, Sequence],
Sequence] = (256, 256, 256),
critic_networks_layer_sizes: Union[Dict[str, Sequence],
Sequence] = (512, 512, 256),
shared_weights: bool = True,
sigma: float = 0.3,
archecture_type: ArchitectureType = ArchitectureType.feedforward,
) -> Mapping[str, types.TensorTransformation]:
"""Default networks for maddpg.
Args:
environment_spec (mava_specs.MAEnvironmentSpec): description of the action and
observation spaces etc. for each agent in the system.
policy_networks_layer_sizes (Union[Dict[str, Sequence], Sequence], optional):
size of policy networks. Defaults to (256, 256, 256).
critic_networks_layer_sizes (Union[Dict[str, Sequence], Sequence], optional):
size of critic networks. Defaults to (512, 512, 256).
shared_weights (bool, optional): whether agents should share weights or not.
Defaults to True.
sigma (float, optional): hyperparameters used to add Gaussian noise for
simple exploration. Defaults to 0.3.
archecture_type (ArchitectureType, optional): archecture used for
agent networks. Can be feedforward or recurrent. Defaults to
ArchitectureType.feedforward.
Returns:
Mapping[str, types.TensorTransformation]: returned agent networks.
"""
# Set Policy function and layer size
if archecture_type == ArchitectureType.feedforward:
policy_network_func = snt.Sequential
elif archecture_type == ArchitectureType.recurrent:
policy_networks_layer_sizes = (128, 128)
policy_network_func = snt.DeepRNN
specs = environment_spec.get_agent_specs()
# Create agent_type specs
if shared_weights:
type_specs = {key.split("_")[0]: specs[key] for key in specs.keys()}
specs = type_specs
if isinstance(policy_networks_layer_sizes, Sequence):
policy_networks_layer_sizes = {
key: policy_networks_layer_sizes
for key in specs.keys()
}
if isinstance(critic_networks_layer_sizes, Sequence):
critic_networks_layer_sizes = {
key: critic_networks_layer_sizes
for key in specs.keys()
}
observation_networks = {}
policy_networks = {}
critic_networks = {}
for key in specs.keys():
# TODO (dries): Make specs[key].actions
# return a list of specs for hybrid action space
# Get total number of action dimensions from action spec.
agent_act_spec = specs[key].actions
if type(specs[key].actions) == DiscreteArray:
num_actions = agent_act_spec.num_values
minimum = [-1.0] * num_actions
maximum = [1.0] * num_actions
agent_act_spec = BoundedArray(
shape=(num_actions, ),
minimum=minimum,
maximum=maximum,
dtype="float32",
name="actions",
)
# Get total number of action dimensions from action spec.
num_dimensions = np.prod(agent_act_spec.shape, dtype=int)
# An optional network to process observations
observation_network = tf2_utils.to_sonnet_module(tf.identity)
# Create the policy network.
if archecture_type == ArchitectureType.feedforward:
policy_network = [
networks.LayerNormMLP(policy_networks_layer_sizes[key],
activate_final=True),
]
elif archecture_type == ArchitectureType.recurrent:
policy_network = [
networks.LayerNormMLP(policy_networks_layer_sizes[key][:-1],
activate_final=True),
snt.LSTM(policy_networks_layer_sizes[key][-1]),
]
policy_network += [
networks.NearZeroInitializedLinear(num_dimensions),
networks.TanhToSpec(agent_act_spec),
]
# Add Gaussian noise for simple exploration.
if sigma and sigma > 0.0:
policy_network += [
networks.ClippedGaussian(sigma),
networks.ClipToSpec(agent_act_spec),
]
policy_network = policy_network_func(policy_network)
# Create the critic network.
critic_network = snt.Sequential([
# The multiplexer concatenates the observations/actions.
networks.CriticMultiplexer(),
networks.LayerNormMLP(list(critic_networks_layer_sizes[key]) + [1],
activate_final=False),
])
observation_networks[key] = observation_network
policy_networks[key] = policy_network
critic_networks[key] = critic_network
return {
"policies": policy_networks,
"critics": critic_networks,
"observations": observation_networks,
}
