Class Evaluator

diffmimic/brax_lib/acting.py:79–153

Class to run evaluations.

# TODO: Consider moving this to its own file.
class Evaluator:
  """Class to run evaluations."""

  def __init__(self, eval_env: envs.Env,
               eval_policy_fn: Callable[[PolicyParams],
                                        Policy],
               eval_encoder_fn: Callable[[PolicyParams],
                                        Policy],
               num_eval_envs: int,
               episode_length: int, action_repeat: int, key: PRNGKey):
    """Init.

    Args:
      eval_env: Batched environment to run evals on.
      eval_policy_fn: Function returning the policy from the policy parameters.
      num_eval_envs: Each env will run 1 episode in parallel for each eval.
      episode_length: Maximum length of an episode.
      action_repeat: Number of physics steps per env step.
      key: RNG key.
    """
    self._key = key
    self._eval_walltime = 0.

    # eval_env = envs.wrappers.EvalWrapper(eval_env)
    eval_env = wrappers.EvalWrapper(eval_env)

    def generate_eval_unroll(cvae_params: PolicyParams,
                             key: PRNGKey,
                             ref_traj: jnp.ndarray,
                             mask: jnp.ndarray) -> (envs.State, brax.QP):
      reset_keys = jax.random.split(key, num_eval_envs)
      # eval_first_state = eval_env.reset(reset_keys)
      eval_first_state = eval_env.reset_ref(reset_keys, ref_traj, mask)
      (normalizer_encoder, normalizer_policy), (encoder_params, policy_params) = cvae_params
      return generate_unroll(
          eval_env,
          eval_first_state,
          eval_policy_fn((normalizer_policy, policy_params)),
          eval_encoder_fn((normalizer_encoder, encoder_params)),
          key,
          unroll_length=episode_length // action_repeat)

    self._generate_eval_unroll = jax.jit(generate_eval_unroll)
    self._steps_per_unroll = episode_length * num_eval_envs

  def run_evaluation(self,
                     cvae_params: PolicyParams,
                     training_metrics: Metrics,
                     ref_traj: jnp.ndarray,
                     mask: jnp.ndarray,
                     aggregate_episodes: bool = True) -> Metrics:
    """Run one epoch of evaluation."""
    self._key, unroll_key = jax.random.split(self._key)

    t = time.time()
    eval_state, (qp_list, latent_list) = self._generate_eval_unroll(cvae_params, unroll_key, ref_traj, mask)
    eval_metrics = eval_state.info['eval_metrics']
    eval_metrics.active_episodes.block_until_ready()