def signature(
cls,
environment_spec: mava_specs.EnvironmentSpec,
extras_spec: tf.TypeSpec = {},
) -> tf.TypeSpec:
# This function currently assumes that self._discount is a scalar.
# If it ever becomes a nested structure and/or a np.ndarray, this method
# will need to know its structure / shape. This is because the signature
# discount shape is the environment's discount shape and this adder's
# discount shape broadcasted together. Also, the reward shape is this
# signature discount shape broadcasted together with the environment
# reward shape. As long as self._discount is a scalar, it will not affect
# either the signature discount shape nor the signature reward shape, so we
# can ignore it.
agent_specs = environment_spec.get_agent_specs()
agents = environment_spec.get_agent_ids()
env_extras_spec = environment_spec.get_extra_specs()
extras_spec.update(env_extras_spec)
obs_specs = {}
act_specs = {}
reward_specs = {}
step_discount_specs = {}
for agent in agents:
rewards_spec, step_discounts_spec = tree_utils.broadcast_structures(
agent_specs[agent].rewards, agent_specs[agent].discounts
)
rewards_spec = tree.map_structure(
_broadcast_specs, rewards_spec, step_discounts_spec
)
step_discounts_spec = tree.map_structure(copy.deepcopy, step_discounts_spec)
obs_specs[agent] = agent_specs[agent].observations
act_specs[agent] = agent_specs[agent].actions
reward_specs[agent] = rewards_spec
step_discount_specs[agent] = step_discounts_spec
transition_spec = [
obs_specs,
act_specs,
extras_spec,
reward_specs,
step_discount_specs,
obs_specs, # next_observation
extras_spec,
]
return tree.map_structure_with_path(
base.spec_like_to_tensor_spec, tuple(transition_spec)
)
def signature(
cls,
environment_spec: specs.EnvironmentSpec,
extras_spec: tf.TypeSpec = {},
) -> tf.TypeSpec:
"""This is a helper method for generating signatures for Reverb tables.
Signatures are useful for validating data types and shapes, see Reverb's
documentation for details on how they are used.
Args:
environment_spec: A `specs.EnvironmentSpec` whose fields are nested
structures with leaf nodes that have `.shape` and `.dtype` attributes.
This should come from the environment that will be used to generate
the data inserted into the Reverb table.
extras_spec: A nested structure with leaf nodes that have `.shape` and
`.dtype` attributes. The structure (and shapes/dtypes) of this must
be the same as the `extras` passed into `ReverbAdder.add`.
Returns:
A `Step` whose leaf nodes are `tf.TensorSpec` objects.
"""
agent_specs = environment_spec.get_agent_specs()
agents = environment_spec.get_agent_ids()
env_extras_spec = environment_spec.get_extra_specs()
extras_spec.update(env_extras_spec)
obs_specs = {}
act_specs = {}
reward_specs = {}
step_discount_specs = {}
for agent in agents:
rewards_spec, step_discounts_spec = tree_utils.broadcast_structures(
agent_specs[agent].rewards, agent_specs[agent].discounts)
obs_specs[agent] = agent_specs[agent].observations
act_specs[agent] = agent_specs[agent].actions
reward_specs[agent] = rewards_spec
step_discount_specs[agent] = step_discounts_spec
spec_step = base.Step(
observations=obs_specs,
actions=act_specs,
rewards=reward_specs,
discounts=step_discount_specs,
start_of_episode=specs.Array(shape=(), dtype=bool),
extras=extras_spec,
)
return tree.map_structure_with_path(base.spec_like_to_tensor_spec,
spec_step)