def _initialize_policy(self):
    """Initialize the policy.

    Run the policy network on dummy data to initialize its parameters for later
    reuse and to analyze the policy distribution. Initializes the attributes
    `self._network` and `self._policy_type`.

    Raises:
      ValueError: Invalid policy distribution.

    Returns:
      Parameters of the policy distribution and policy state.
    """
    with tf.device('/gpu:0' if self._use_gpu else '/cpu:0'):
      network = functools.partial(
          self._config.network, self._config, self._batch_env.action_space)
      self._network = tf.make_template('network', network)
      output = self._network(
          tf.zeros_like(self._batch_env.observ)[:, None],
          tf.ones(len(self._batch_env)))
    if output.policy.event_shape != self._batch_env.action.shape[1:]:
      message = 'Policy event shape {} does not match action shape {}.'
      message = message.format(
          output.policy.event_shape, self._batch_env.action.shape[1:])
      raise ValueError(message)
    self._policy_type = type(output.policy)
    is_tensor = lambda x: isinstance(x, tf.Tensor)
    policy_params = tools.nested.filter(is_tensor, output.policy.parameters)
    set_batch_dim = lambda x: utility.set_dimension(x, 0, len(self._batch_env))
    tools.nested.map(set_batch_dim, policy_params)
    if output.state is not None:
      tools.nested.map(set_batch_dim, output.state)
    return policy_params, output.state
Initialize the policy.

    Run the policy network on dummy data to initialize its parameters for later
    reuse and to analyze the policy distribution. Initializes the attributes
    `self._network` and `self._policy_type`.

    Raises:
      ValueError: Invalid policy distribution.

    Returns:
      Parameters of the policy distribution and policy state.