def default_eval_actor(
source: core.VariableSource) -> types.ModelToSnapshot:
params = source.get_variables(['policy'])[0]
return types.ModelToSnapshot(apply_policy_and_sample(networks, True),
params, {
'key': dummy_key,
'obs': dummy_obs
})
def _get_model1(variable_source: core.VariableSource) -> types.ModelToSnapshot:
return types.ModelToSnapshot(
model=_model1,
params=variable_source.get_variables(['params_model1'])[0],
dummy_kwargs={
'x': jnp.ones([3, 1], dtype=jnp.float32),
},
)
def q_network(source: core.VariableSource) -> types.ModelToSnapshot:
params = source.get_variables(['critic'])[0]
return types.ModelToSnapshot(networks.q_network.apply, params, {
'obs': dummy_obs,
'action': dummy_action
})