def output_signature(self, *, function):
if function == 'action_value':
return SignatureDict(
singleton=TensorSpec(type='float', shape=()).signature(
batched=True))
if function == 'action_values':
return SignatureDict(
singleton=self.actions_spec.fmap(function=(
lambda spec: TensorSpec(type='float', shape=spec.shape
).signature(batched=True)),
cls=SignatureDict))
elif function == 'state_values':
return SignatureDict(
singleton=self.actions_spec.fmap(function=(
lambda spec: TensorSpec(type='float', shape=spec.shape
).signature(batched=True)),
cls=SignatureDict))
else:
try:
return Policy.output_signature(self=self, function=function)
except NotImplementedError:
try:
return StateValue.output_signature(self=self,
function=function)
except NotImplementedError:
return ActionValue.output_signature(self=self,
function=function)
def input_signature(self, *, function):
if function == 'apply':
assert len(self.internals_spec
) == 0 or self.temporal_processing == 'iterative'
return SignatureDict(
x=self.input_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True))
elif function == 'cumulative_apply':
assert self.temporal_processing == 'cumulative'
cumulative_input_spec = self.input_spec.copy()
cumulative_input_spec.shape = (
None, ) + cumulative_input_spec.shape
return SignatureDict(
xs=cumulative_input_spec.signature(batched=True),
lengths=TensorSpec(type='int',
shape=()).signature(batched=True))
elif function == 'iterative_apply':
assert self.temporal_processing == 'iterative'
return SignatureDict(
x=self.input_spec.signature(batched=True),
internals=self.internals_spec.signature(batched=True))
elif function == 'past_horizon':
return SignatureDict()
else:
return super().input_signature(function=function)
def output_signature(self, *, function):
if function == 'loss':
return SignatureDict(
singleton=TensorSpec(type='float', shape=()).signature(batched=True)
)
elif function == 'reference':
return SignatureDict(singleton=self.reference_spec().signature(batched=True))
else:
return super().output_signature(function=function)
def output_signature(self, *, function):
if function == 'apply':
return SignatureDict(singleton=self.output_spec().signature(
batched=True))
elif function == 'reset':
return SignatureDict(
singleton=TensorSpec(type='bool', shape=()).signature(
batched=False))
else:
return super().output_signature(function=function)
def output_signature(self, *, function):
if function == 'next_internals':
return SignatureDict(singleton=self.internals_spec.signature(
batched=True))
elif function == 'past_horizon':
return SignatureDict(
singleton=TensorSpec(type='int', shape=()).signature(
batched=False))
else:
return super().output_signature(function=function)
def input_signature(self, *, function):
if function == 'apply':
return SignatureDict(
x=self.inputs_spec.signature(batched=True),
deterministic=TensorSpec(type='bool',
shape=()).signature(batched=False))
elif function == 'reset':
return SignatureDict()
else:
return super().input_signature(function=function)
def input_signature(self, *, function):
if function == 'apply':
return SignatureDict(
x=self.inputs_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True))
elif function == 'past_horizon':
return SignatureDict()
else:
return super().input_signature(function=function)
def output_signature(self, *, function):
if function == 'apply':
return SignatureDict(singleton=self.output_spec().signature(
batched=True))
else:
return super().output_signature(function=function)
def input_signature(self, *, function):
if function == 'next_internals':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
actions=self.actions_spec.signature(batched=True),
deterministic=TensorSpec(type='bool',
shape=()).signature(batched=False))
elif function == 'past_horizon':
return SignatureDict()
else:
return super().input_signature(function=function)
def input_signature(self, *, function):
if function == 'step' or function == 'update':
return SignatureDict(arguments=self.arguments_spec.signature(
batched=True))
else:
return super().input_signature(function=function)
def output_signature(self, *, function):
if function == 'regularize':
return SignatureDict(
singleton=TensorSpec(type='float', shape=()).signature(
batched=False))
else:
raise NotImplementedError
def output_signature(self, *, function):
if function == 'state_value':
return SignatureDict(
singleton=TensorSpec(type='float', shape=()).signature(batched=True)
)
else:
return super().output_signature(function=function)
def output_signature(self, *, function):
if function == 'act':
return SignatureDict(
actions=self.actions_spec.signature(batched=True),
internals=self.internals_spec.signature(batched=True))
else:
return super().output_signature(function=function)
def input_signature(self, *, function):
if function == 'actions_value':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True),
actions=self.actions_spec.signature(batched=True))
elif function == 'actions_values':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True),
actions=self.actions_spec.signature(batched=True))
elif function == 'all_actions_values':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True))
elif function == 'states_value':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True))
elif function == 'states_values':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True))
else:
return super().input_signature(function=function)
def input_signature(self, *, function):
if function == 'end' or function == 'next_step' or function == 'step':
return SignatureDict(
arguments=self.arguments_spec.signature(batched=True),
x=self.values_spec.signature(batched=False),
conjugate=self.values_spec.signature(batched=False),
residual=self.values_spec.signature(batched=False),
squared_residual=TensorSpec(type='float',
shape=()).signature(batched=False))
elif function == 'solve' or function == 'start':
return SignatureDict(
arguments=self.arguments_spec.signature(batched=True),
x_init=self.values_spec.signature(batched=False),
b=self.values_spec.signature(batched=False))
else:
return super().input_signature(function=function)
def input_signature(self, *, function):
if function == 'apply':
return self.inputs_spec.signature(batched=True)
elif function == 'reset':
return SignatureDict()
else:
return super().input_signature(function=function)
def input_signature(self, *, function):
if function == 'act':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int', shape=(2,)).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True)
)
elif function == 'join_value_per_action':
return SignatureDict(
values=self.actions_spec.fmap(
function=(lambda x: TensorSpec(type='float', shape=x.shape))
).signature(batched=True)
)
elif function == 'past_horizon':
return SignatureDict()
else:
return super().input_signature(function=function)
def input_signature(self, *, function):
if function == 'enqueue':
return self.values_spec.signature(batched=True)
elif function == 'predecessors':
return SignatureDict(
indices=TensorSpec(type='int', shape=()).signature(batched=True),
horizon=TensorSpec(type='int', shape=()).signature(batched=False)
)
elif function == 'reset':
return SignatureDict()
elif function == 'retrieve':
return SignatureDict(indices=TensorSpec(type='int', shape=()).signature(batched=True))
elif function == 'retrieve_episodes':
return SignatureDict(n=TensorSpec(type='int', shape=()).signature(batched=False))
elif function == 'retrieve_timesteps':
return SignatureDict(
n=TensorSpec(type='int', shape=()).signature(batched=False),
past_horizon=TensorSpec(type='int', shape=()).signature(batched=False),
future_horizon=TensorSpec(type='int', shape=()).signature(batched=False)
)
elif function == 'successors':
return SignatureDict(
indices=TensorSpec(type='int', shape=()).signature(batched=True),
horizon=TensorSpec(type='int', shape=()).signature(batched=False)
)
else:
return super().input_signature(function=function)
def output_signature(self, *, function):
if function == 'act_entropy':
return SignatureDict(
actions=self.actions_spec.signature(batched=True),
internals=self.internals_spec.signature(batched=True),
entropy=TensorSpec(type='float',
shape=()).signature(batched=True))
elif function == 'entropy':
return SignatureDict(
singleton=TensorSpec(type='float', shape=()).signature(
batched=True))
elif function == 'entropies':
return SignatureDict(
singleton=self.actions_spec.fmap(function=(
lambda spec: TensorSpec(type='float', shape=spec.shape
).signature(batched=True)),
cls=SignatureDict))
elif function == 'kl_divergence':
return SignatureDict(
singleton=TensorSpec(type='float', shape=()).signature(
batched=True))
elif function == 'kl_divergences':
return SignatureDict(
singleton=self.actions_spec.fmap(function=(
lambda spec: TensorSpec(type='float', shape=spec.shape
).signature(batched=True)),
cls=SignatureDict))
elif function == 'kldiv_reference':
return SignatureDict(singleton=self.kldiv_reference_spec.signature(
batched=True))
elif function == 'log_probability':
return SignatureDict(
singleton=TensorSpec(type='float', shape=()).signature(
batched=True))
elif function == 'log_probabilities':
return SignatureDict(
singleton=self.actions_spec.fmap(function=(
lambda spec: TensorSpec(type='float', shape=spec.shape
).signature(batched=True)),
cls=SignatureDict))
else:
return super().output_signature(function=function)
def input_signature(self, *, function):
if function == 'action_value':
return SignatureDict(
parameters=self.parameters_spec.signature(batched=True),
action=self.action_spec.signature(batched=True))
elif function == 'all_action_values':
return SignatureDict(parameters=self.parameters_spec.signature(
batched=True))
elif function == 'entropy':
return SignatureDict(parameters=self.parameters_spec.signature(
batched=True))
elif function == 'kl_divergence':
return SignatureDict(
parameters1=self.parameters_spec.signature(batched=True),
parameters2=self.parameters_spec.signature(batched=True))
elif function == 'log_probability':
return SignatureDict(
parameters=self.parameters_spec.signature(batched=True),
action=self.action_spec.signature(batched=True))
elif function == 'mode':
return SignatureDict(parameters=self.parameters_spec.signature(
batched=True))
elif function == 'parametrize':
return SignatureDict(
x=self.input_spec.signature(batched=True),
conditions=self.conditions_spec.signature(batched=True))
elif function == 'sample':
return SignatureDict(
parameters=self.parameters_spec.signature(batched=True),
temperature=TensorSpec(type='float',
shape=()).signature(batched=False))
elif function == 'states_value':
return SignatureDict(parameters=self.parameters_spec.signature(
batched=True))
else:
return super().input_signature(function=function)
def get_output_signature(sequence_values, initial_values):
if len(sequence_values) == 0:
if len(initial_values) == 0:
return SignatureDict(
singleton=TensorSpec(type='int', shape=()).signature(batched=True)
)
else:
return SignatureDict(
lengths=TensorSpec(type='int', shape=()).signature(batched=True),
initial_values=self.values_spec[initial_values].signature(batched=True)
)
elif len(initial_values) == 0:
return SignatureDict(
starts_lengths=TensorSpec(type='int', shape=(2,)).signature(batched=True),
sequence_values=self.values_spec[sequence_values].signature(batched=True)
)
else:
return SignatureDict(
starts_lengths=TensorSpec(type='int', shape=(2,)).signature(batched=True),
sequence_values=self.values_spec[sequence_values].signature(batched=True),
initial_values=self.values_spec[initial_values].signature(batched=True)
)
def output_signature(self, *, function):
if function == 'step':
return self.variables_spec.fmap(
function=(lambda spec: spec.signature(batched=False)),
cls=SignatureDict)
elif function == 'update':
return SignatureDict(
singleton=TensorSpec(type='bool', shape=()).signature(
batched=False))
else:
return super().output_signature(function=function)
def input_signature(self, *, function):
if function == 'end' or function == 'next_step' or function == 'step':
return SignatureDict(
arguments=self.arguments_spec.signature(batched=True),
x=self.values_spec.signature(batched=False),
deltas=self.values_spec.signature(batched=False),
improvement=TensorSpec(type='float', shape=()).signature(batched=False),
last_improvement=TensorSpec(type='float', shape=()).signature(batched=False),
base_value=TensorSpec(type='float', shape=()).signature(batched=False),
estimated=TensorSpec(type='float', shape=()).signature(batched=False)
)
elif function == 'solve' or function == 'start':
return SignatureDict(
arguments=self.arguments_spec.signature(batched=True),
x_init=self.values_spec.signature(batched=False),
base_value=TensorSpec(type='float', shape=()).signature(batched=False),
zero_value=TensorSpec(type='float', shape=()).signature(batched=False),
estimated=TensorSpec(type='float', shape=()).signature(batched=False)
)
else:
return super().input_signature(function=function)
def input_signature(self, *, function):
if function == 'loss':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int', shape=(2,)).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True),
actions=self.actions_spec.signature(batched=True),
reward=self.reward_spec.signature(batched=True),
reference=self.reference_spec().signature(batched=True)
)
elif function == 'reference':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int', shape=(2,)).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True),
actions=self.actions_spec.signature(batched=True)
)
else:
return super().input_signature(function=function)
def input_signature(self, *, function):
if function == 'action_value':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True),
actions=self.actions_spec.signature(batched=True))
elif function == 'action_values':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True),
actions=self.actions_spec.signature(batched=True))
elif function == 'state_values':
return SignatureDict(
states=self.states_spec.signature(batched=True),
horizons=TensorSpec(type='int',
shape=(2, )).signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True))
else:
try:
return Policy.input_signature(self=self, function=function)
except NotImplementedError:
try:
return StateValue.input_signature(self=self,
function=function)
except NotImplementedError:
return ActionValue.input_signature(self=self,
function=function)
def output_signature(self, *, function):
if function == 'end' or function == 'solve':
return SignatureDict(singleton=self.values_spec.signature(
batched=False))
elif function == 'next_step':
return SignatureDict(
singleton=TensorSpec(type='bool', shape=()).signature(
batched=False))
elif function == 'start' or function == 'step':
return SignatureDict(
arguments=self.arguments_spec.signature(batched=True),
x=self.values_spec.signature(batched=False),
deltas=self.values_spec.signature(batched=False),
improvement=TensorSpec(type='float',
shape=()).signature(batched=False),
last_improvement=TensorSpec(type='float',
shape=()).signature(batched=False),
base_value=TensorSpec(type='float',
shape=()).signature(batched=False))
else:
return super().output_signature(function=function)
def output_signature(self, *, function):
shape = self.action_spec.shape
if function == 'action_value':
return SignatureDict(
singleton=TensorSpec(type='float', shape=shape).signature(
batched=True))
elif function == 'entropy':
return SignatureDict(
singleton=TensorSpec(type='float', shape=shape).signature(
batched=True))
elif function == 'kl_divergence':
return SignatureDict(
singleton=TensorSpec(type='float', shape=shape).signature(
batched=True))
elif function == 'log_probability':
return SignatureDict(
singleton=TensorSpec(type='float', shape=shape).signature(
batched=True))
elif function == 'mode':
return SignatureDict(singleton=self.action_spec.signature(
batched=True))
elif function == 'parametrize':
return SignatureDict(singleton=self.parameters_spec.signature(
batched=True))
elif function == 'sample':
return SignatureDict(singleton=self.action_spec.signature(
batched=True))
elif function == 'state_value':
return SignatureDict(
singleton=TensorSpec(type='float', shape=shape).signature(
batched=True))
else:
return super().output_signature(function=function)
def output_signature(self, *, function):
if function == 'act':
return SignatureDict(
actions=self.actions_spec.signature(batched=True),
timesteps=TensorSpec(type='int',
shape=()).signature(batched=False))
if function == 'core_act':
return SignatureDict(
actions=self.actions_spec.signature(batched=True),
internals=self.internals_spec.signature(batched=True))
elif function == 'core_observe':
return SignatureDict(
singleton=TensorSpec(type='bool', shape=()).signature(
batched=False))
elif function == 'independent_act':
if len(self.internals_spec) > 0:
return SignatureDict(
actions=self.actions_spec.signature(batched=True),
internals=self.internals_spec.signature(batched=True))
else:
return SignatureDict(singleton=self.actions_spec.signature(
batched=True))
elif function == 'observe':
return SignatureDict(
updated=TensorSpec(type='bool',
shape=()).signature(batched=False),
episodes=TensorSpec(type='int',
shape=()).signature(batched=False),
updates=TensorSpec(type='int',
shape=()).signature(batched=False))
elif function == 'reset':
return SignatureDict(
timesteps=TensorSpec(type='int',
shape=()).signature(batched=False),
episodes=TensorSpec(type='int',
shape=()).signature(batched=False),
updates=TensorSpec(type='int',
shape=()).signature(batched=False))
else:
return super().output_signature(function=function)
def input_signature(self, *, function):
if function == 'act':
return SignatureDict(
states=self.states_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True),
parallel=self.parallel_spec.signature(batched=True))
elif function == 'core_act':
return SignatureDict(
states=self.states_spec.signature(batched=True),
internals=self.internals_spec.signature(batched=True),
auxiliaries=self.auxiliaries_spec.signature(batched=True),
parallel=self.parallel_spec.signature(batched=True),
deterministic=self.deterministic_spec.signature(batched=False))
elif function == 'core_observe':
return SignatureDict(
terminal=self.terminal_spec.signature(batched=True),
reward=self.reward_spec.signature(batched=True),
parallel=self.parallel_spec.signature(batched=False))
elif function == 'independent_act':
signature = SignatureDict(states=self.states_spec.signature(
batched=True))
if len(self.internals_spec) > 0:
signature['internals'] = self.internals_spec.signature(
batched=True)
if len(self.auxiliaries_spec) > 0:
signature['auxiliaries'] = self.auxiliaries_spec.signature(
batched=True)
signature['deterministic'] = self.deterministic_spec.signature(
batched=False)
return signature
elif function == 'observe':
return SignatureDict(
terminal=self.terminal_spec.signature(batched=True),
reward=self.reward_spec.signature(batched=True),
parallel=self.parallel_spec.signature(batched=False))
elif function == 'reset':
return SignatureDict()
else:
return super().input_signature(function=function)
def output_signature(self, *, function):
if function == 'apply':
if self.temporal_processing == 'cumulative':
return SignatureDict(singleton=self.output_spec().signature(
batched=True))
elif self.temporal_processing == 'iterative':
return SignatureDict(
x=self.output_spec().signature(batched=True),
internals=self.internals_spec.signature(batched=True))
elif function == 'cumulative_apply':
assert self.temporal_processing == 'cumulative'
return SignatureDict(singleton=self.output_spec().signature(
batched=True))
elif function == 'iterative_apply':
assert self.temporal_processing == 'iterative'
return SignatureDict(
x=self.output_spec().signature(batched=True),
internals=self.internals_spec.signature(batched=True))
elif function == 'iterative_body':
assert self.temporal_processing == 'iterative'
return SignatureDict(
x=self.input_spec.signature(batched=True),
indices=TensorSpec(type='int',
shape=()).signature(batched=True),
remaining=TensorSpec(type='int',
shape=()).signature(batched=True),
current_x=self.output_spec().signature(batched=True),
current_internals=self.internals_spec.signature(batched=True))
elif function == 'past_horizon':
return SignatureDict(
singleton=TensorSpec(type='int', shape=()).signature(
batched=False))
else:
return super().output_signature(function=function)