def load_policy(policy_cfg, action_spec):
"""Loads policy based on config."""
if policy_cfg.ptype not in PTYPES:
raise ValueError('Unknown policy type %s.' % policy_cfg.ptype)
if policy_cfg.ptype == 'randwalk':
policy = policies.ContinuousRandomPolicy(action_spec)
elif policy_cfg.ptype in ['randinit', 'load']:
a_net = networks.ActorNetwork(action_spec,
fc_layer_params=policy_cfg.model_params)
if policy_cfg.ptype == 'load':
logging.info('Loading policy from %s...', policy_cfg.ckpt)
policy_ckpt = tf.train.Checkpoint(policy=a_net)
policy_ckpt.restore(policy_cfg.ckpt).expect_partial()
policy = wrap_policy(a_net, policy_cfg.wrapper)
return policy
def p_net_factory():
return networks.ActorNetwork(
action_spec,
fc_layer_params=model_params[0])