code: import numpy as np
  import tensorflow as tf
  
  from garage.core import Serializable
  from garage.misc import special
  from garage.misc.overrides import overrides
  from garage.tf.core import LayersPowered
  import garage.tf.core.layers as L
  from garage.tf.core.network import GRUNetwork
  from garage.tf.distributions import RecurrentCategorical
  from garage.tf.misc import tensor_utils
  from garage.tf.policies import StochasticPolicy
  from garage.tf.spaces import Discrete
  
  
  class CategoricalGRUPolicy(StochasticPolicy, LayersPowered, Serializable):
      def __init__(
              self,
              env_spec,
              name="CategoricalGRUPolicy",
              hidden_dim=32,
              feature_network=None,
              state_include_action=True,
              hidden_nonlinearity=tf.tanh,
              gru_layer_cls=L.GRULayer,
      ):
          """
          :param env_spec: A spec for the env.
          :param hidden_dim: dimension of hidden layer
          :param hidden_nonlinearity: nonlinearity used for each hidden layer
          :return:
          """
          assert isinstance(env_spec.action_space, Discrete)
  
          self._prob_network_name = "prob_network"
          with tf.variable_scope(name, "CategoricalGRUPolicy"):
              Serializable.quick_init(self, locals())
              super(CategoricalGRUPolicy, self).__init__(env_spec)
  
              obs_dim = env_spec.observation_space.flat_dim
              action_dim = env_spec.action_space.flat_dim
  
              if state_include_action:
                  input_dim = obs_dim + action_dim
              else:
                  input_dim = obs_dim
  
              l_input = L.InputLayer(shape=(None, None, input_dim), name="input")
  
              if feature_network is None:
                  feature_dim = input_dim
                  l_flat_feature = None
                  l_feature = l_input
              else:
                  feature_dim = feature_network.output_layer.output_shape[-1]
                  l_flat_feature = feature_network.output_layer
                  l_feature = L.OpLayer(
                      l_flat_feature,
                      extras=[l_input],
                      name="reshape_feature",
                      op=lambda flat_feature, input: tf.reshape(
                          flat_feature,
                          tf.stack([
                              tf.shape(input)[0],
                              tf.shape(input)[1], feature_dim
                          ])),
                      shape_op=lambda _, input_shape: (
                          input_shape[0], input_shape[1], feature_dim))
  
              prob_network = GRUNetwork(
                  input_shape=(feature_dim, ),
                  input_layer=l_feature,
                  output_dim=env_spec.action_space.n,
                  hidden_dim=hidden_dim,
                  hidden_nonlinearity=hidden_nonlinearity,
                  output_nonlinearity=tf.nn.softmax,
                  gru_layer_cls=gru_layer_cls,
                  name=self._prob_network_name)
  
              self.prob_network = prob_network
              self.feature_network = feature_network
              self.l_input = l_input
              self.state_include_action = state_include_action
  
              flat_input_var = tf.placeholder(
                  dtype=tf.float32, shape=(None, input_dim), name="flat_input")
              if feature_network is None:
                  feature_var = flat_input_var
              else:
                  with tf.name_scope("feature_network", values=[flat_input_var]):
                      feature_var = L.get_output(
                          l_flat_feature,
                          {feature_network.input_layer: flat_input_var})
  
              with tf.name_scope(self._prob_network_name, values=[feature_var]):
                  out_prob_step, out_prob_hidden = L.get_output(
                      [
                          prob_network.step_output_layer,
                          prob_network.step_hidden_layer
                      ], {prob_network.step_input_layer: feature_var})
                  out_prob_step = tf.identity(out_prob_step, "prob_step_output")
                  out_prob_hidden = tf.identity(out_prob_hidden,
                                                "prob_step_hidden")
  
              self.f_step_prob = tensor_utils.compile_function(
                  [flat_input_var, prob_network.step_prev_state_layer.input_var],
                  [out_prob_step, out_prob_hidden])
  
              self.input_dim = input_dim
              self.action_dim = action_dim
              self.hidden_dim = hidden_dim
              self.name = name
  
              self.prev_actions = None
              self.prev_hiddens = None
              self.dist = RecurrentCategorical(env_spec.action_space.n)
  
              out_layers = [prob_network.output_layer]
              if feature_network is not None:
                  out_layers.append(feature_network.output_layer)
  
              LayersPowered.__init__(self, out_layers)
  
      @overrides
      def dist_info_sym(self, obs_var, state_info_vars, name=None):
          with tf.name_scope(name, "dist_info_sym", [obs_var, state_info_vars]):
              n_batches = tf.shape(obs_var)[0]
              n_steps = tf.shape(obs_var)[1]
              obs_var = tf.reshape(obs_var, tf.stack([n_batches, n_steps, -1]))
              obs_var = tf.cast(obs_var, tf.float32)
              if self.state_include_action:
                  prev_action_var = tf.cast(state_info_vars["prev_action"],
                                            tf.float32)
                  all_input_var = tf.concat(
                      axis=2, values=[obs_var, prev_action_var])
              else:
                  all_input_var = obs_var
              if self.feature_network is None:
                  with tf.name_scope(
                          self._prob_network_name, values=[all_input_var]):
                      prob = L.get_output(self.prob_network.output_layer,
                                          {self.l_input: all_input_var})
                  return dict(prob=prob)
              else:
                  flat_input_var = tf.reshape(all_input_var,
                                              (-1, self.input_dim))
                  with tf.name_scope(
                          self._prob_network_name,
                          values=[all_input_var, flat_input_var]):
                      prob = L.get_output(
                          self.prob_network.output_layer, {
                              self.l_input: all_input_var,
                              self.feature_network.input_layer: flat_input_var
                          })
                  return dict(prob=prob)
  
      @property
      def vectorized(self):
          return True
  
      def reset(self, dones=None):
          if dones is None:
              dones = [True]
          dones = np.asarray(dones)
          if self.prev_actions is None or len(dones) != len(self.prev_actions):
              self.prev_actions = np.zeros((len(dones),
                                            self.action_space.flat_dim))
              self.prev_hiddens = np.zeros((len(dones), self.hidden_dim))
  
          self.prev_actions[dones] = 0.
          self.prev_hiddens[
              dones] = self.prob_network.hid_init_param.eval()  # get_value()
  
      # The return value is a pair. The first item is a matrix (N, A), where each
      # entry corresponds to the action value taken. The second item is a vector
      # of length N, where each entry is the density value for that action, under
      # the current policy
      @overrides
      def get_action(self, observation):
          actions, agent_infos = self.get_actions([observation])
          return actions[0], {k: v[0] for k, v in agent_infos.items()}
  
      @overrides
      def get_actions(self, observations):
          flat_obs = self.observation_space.flatten_n(observations)
          if self.state_include_action:
              assert self.prev_actions is not None
              all_input = np.concatenate([flat_obs, self.prev_actions], axis=-1)
          else:
              all_input = flat_obs
          probs, hidden_vec = self.f_step_prob(all_input, self.prev_hiddens)
          actions = special.weighted_sample_n(probs,
                                              np.arange(self.action_space.n))
          prev_actions = self.prev_actions
          self.prev_actions = self.action_space.flatten_n(actions)
          self.prev_hiddens = hidden_vec
          agent_info = dict(prob=probs)
          if self.state_include_action:
              agent_info["prev_action"] = np.copy(prev_actions)
          return actions, agent_info
  
      @property
      @overrides
      def recurrent(self):
          return True
  
      @property
      def distribution(self):
          return self.dist
  
      @property
      def state_info_specs(self):
          if self.state_include_action:
              return [
                  ("prev_action", (self.action_dim, )),
              ]
          else:
              return []
  
apis: tensorflow.concat tensorflow.shape numpy.asarray numpy.arange tensorflow.stack tensorflow.cast tensorflow.identity tensorflow.placeholder tensorflow.reshape numpy.concatenate numpy.copy tensorflow.name_scope tensorflow.variable_scope
api_extract: [(164, 'numpy.asarray', 'np.asarray', (['dones'], {}), True, 'import numpy as np\n'), (36, 'tensorflow.variable_scope', 'tf.variable_scope', (['name', '"""CategoricalGRUPolicy"""'], {}), True, 'import tensorflow as tf\n'), (48, 'garage.tf.core.layers.InputLayer', 'L.InputLayer', ([], {'shape': '(None, None, input_dim)', 'name': '"""input"""'}), True, 'import garage.tf.core.layers as L\n'), (70, 'garage.tf.core.network.GRUNetwork', 'GRUNetwork', ([], {'input_shape': '(feature_dim,)', 'input_layer': 'l_feature', 'output_dim': 'env_spec.action_space.n', 'hidden_dim': 'hidden_dim', 'hidden_nonlinearity': 'hidden_nonlinearity', 'output_nonlinearity': 'tf.nn.softmax', 'gru_layer_cls': 'gru_layer_cls', 'name': 'self._prob_network_name'}), False, 'from garage.tf.core.network import GRUNetwork\n'), (85, 'tensorflow.placeholder', 'tf.placeholder', ([], {'dtype': 'tf.float32', 'shape': '(None, input_dim)', 'name': '"""flat_input"""'}), True, 'import tensorflow as tf\n'), (105, 'garage.tf.misc.tensor_utils.compile_function', 'tensor_utils.compile_function', (['[flat_input_var, prob_network.step_prev_state_layer.input_var]', '[out_prob_step, out_prob_hidden]'], {}), False, 'from garage.tf.misc import tensor_utils\n'), (116, 'garage.tf.distributions.RecurrentCategorical', 'RecurrentCategorical', (['env_spec.action_space.n'], {}), False, 'from garage.tf.distributions import RecurrentCategorical\n'), (122, 'garage.tf.core.LayersPowered.__init__', 'LayersPowered.__init__', (['self', 'out_layers'], {}), False, 'from garage.tf.core import LayersPowered\n'), (126, 'tensorflow.name_scope', 'tf.name_scope', (['name', '"""dist_info_sym"""', '[obs_var, state_info_vars]'], {}), True, 'import tensorflow as tf\n'), (130, 'tensorflow.cast', 'tf.cast', (['obs_var', 'tf.float32'], {}), True, 'import tensorflow as tf\n'), (188, 'numpy.concatenate', 'np.concatenate', (['[flat_obs, self.prev_actions]'], {'axis': '(-1)'}), True, 'import numpy as np\n'), (193, 'numpy.arange', 'np.arange', (['self.action_space.n'], {}), True, 'import numpy as np\n'), (199, 'numpy.copy', 'np.copy', (['prev_actions'], {}), True, 'import numpy as np\n'), (95, 'tensorflow.name_scope', 'tf.name_scope', (['self._prob_network_name'], {'values': '[feature_var]'}), True, 'import tensorflow as tf\n'), (96, 'garage.tf.core.layers.get_output', 'L.get_output', (['[prob_network.step_output_layer, prob_network.step_hidden_layer]', '{prob_network.step_input_layer: feature_var}'], {}), True, 'import garage.tf.core.layers as L\n'), (101, 'tensorflow.identity', 'tf.identity', (['out_prob_step', '"""prob_step_output"""'], {}), True, 'import tensorflow as tf\n'), (102, 'tensorflow.identity', 'tf.identity', (['out_prob_hidden', '"""prob_step_hidden"""'], {}), True, 'import tensorflow as tf\n'), (127, 'tensorflow.shape', 'tf.shape', (['obs_var'], {}), True, 'import tensorflow as tf\n'), (128, 'tensorflow.shape', 'tf.shape', (['obs_var'], {}), True, 'import tensorflow as tf\n'), (129, 'tensorflow.stack', 'tf.stack', (['[n_batches, n_steps, -1]'], {}), True, 'import tensorflow as tf\n'), (132, 'tensorflow.cast', 'tf.cast', (["state_info_vars['prev_action']", 'tf.float32'], {}), True, 'import tensorflow as tf\n'), (134, 'tensorflow.concat', 'tf.concat', ([], {'axis': '(2)', 'values': '[obs_var, prev_action_var]'}), True, 'import tensorflow as tf\n'), (145, 'tensorflow.reshape', 'tf.reshape', (['all_input_var', '(-1, self.input_dim)'], {}), True, 'import tensorflow as tf\n'), (90, 'tensorflow.name_scope', 'tf.name_scope', (['"""feature_network"""'], {'values': '[flat_input_var]'}), True, 'import tensorflow as tf\n'), (91, 'garage.tf.core.layers.get_output', 'L.get_output', (['l_flat_feature', '{feature_network.input_layer: flat_input_var}'], {}), True, 'import garage.tf.core.layers as L\n'), (139, 'tensorflow.name_scope', 'tf.name_scope', (['self._prob_network_name'], {'values': '[all_input_var]'}), True, 'import tensorflow as tf\n'), (141, 'garage.tf.core.layers.get_output', 'L.get_output', (['self.prob_network.output_layer', '{self.l_input: all_input_var}'], {}), True, 'import garage.tf.core.layers as L\n'), (147, 'tensorflow.name_scope', 'tf.name_scope', (['self._prob_network_name'], {'values': '[all_input_var, flat_input_var]'}), True, 'import tensorflow as tf\n'), (150, 'garage.tf.core.layers.get_output', 'L.get_output', (['self.prob_network.output_layer', '{self.l_input: all_input_var, self.feature_network.input_layer: flat_input_var}'], {}), True, 'import garage.tf.core.layers as L\n'), (64, 'tensorflow.shape', 'tf.shape', (['input'], {}), True, 'import tensorflow as tf\n'), (65, 'tensorflow.shape', 'tf.shape', (['input'], {}), True, 'import tensorflow as tf\n')]