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])