| 68 | make_encoder = encoder_networks.make_inference_fn(encoder_network) |
| 69 | |
| 70 | def make_evaluator(env): |
| 71 | |
| 72 | ep_len_eval = 196 |
| 73 | |
| 74 | eval_env = wrappers.EpisodeWrapper(env, ep_len_eval - 1, 1) |
| 75 | eval_env = wrappers.VmapWrapper(eval_env) |
| 76 | eval_env = wrappers.AutoResetWrapper(eval_env) |
| 77 | return acting.Evaluator( |
| 78 | eval_env, |
| 79 | eval_policy_fn=functools.partial(make_policy, deterministic=False), |
| 80 | eval_encoder_fn=functools.partial(make_encoder, deterministic=False), |
| 81 | num_eval_envs=4, |
| 82 | episode_length=ep_len_eval, |
| 83 | action_repeat=1, |
| 84 | key=jax.random.PRNGKey(999) |
| 85 | ) |
| 86 | |
| 87 | evaluator_global = make_evaluator(env_global) |
| 88 | evaluator_global_perturb = make_evaluator(env_global_perturb) |