code: def _initialize_policy(self):
      
      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
docs: 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.