docstr: Test a policy that involves recurrent layers.
code: def test_recurrent_states(self):
      """Test a policy that involves recurrent layers."""
  
      # The environment just has a constant state.
  
      class TestEnvironment(dc.rl.Environment):
          def __init__(self):
              super(TestEnvironment, self).__init__((10,), 10)
              self._state = np.random.random(10).astype(np.float32)
  
          def step(self, action):
              self._state = np.random.random(10).astype(np.float32)
              return 0.0
  
          def reset(self):
              pass
  
      # The policy includes a single recurrent layer.
  
      class TestPolicy(dc.rl.Policy):
          def __init__(self):
              super(TestPolicy, self).__init__(
                  ["action_prob", "value", "rnn_state"], [np.zeros(10)]
              )
  
          def create_model(self, **kwargs):
              state = Input(shape=(10,))
              rnn_state = Input(shape=(10,))
              reshaped = Reshape((1, 10))(state)
              gru, rnn_final_state = GRU(
                  10,
                  return_state=True,
                  return_sequences=True,
                  time_major=True,
              )(reshaped, initial_state=rnn_state)
              output = Softmax()(Reshape((10,))(gru))
              value = dc.models.layers.Variable([0.0])([state])
              return tf.keras.Model(
                  inputs=[state, rnn_state],
                  outputs=[output, value, rnn_final_state],
              )
  
      # We don't care about actually optimizing it, so just run a few rollouts to make
      # sure fit() doesn't crash, then check the behavior of the GRU state.
  
      env = TestEnvironment()
      a2c = dc.rl.A2C(env, TestPolicy())
      a2c.fit(100)
      # On the first call, the initial state should be all zeros.
      prob1, value1 = a2c.predict(
          env.state, use_saved_states=True, save_states=False
      )
      # It should still be zeros since we didn't save it last time.
      prob2, value2 = a2c.predict(
          env.state, use_saved_states=True, save_states=True
      )
      # It should be different now.
      prob3, value3 = a2c.predict(
          env.state, use_saved_states=True, save_states=False
      )
      # This should be the same as the previous one.
      prob4, value4 = a2c.predict(
          env.state, use_saved_states=True, save_states=False
      )
      # Now we reset it, so we should get the same result as initially.
      prob5, value5 = a2c.predict(
          env.state, use_saved_states=False, save_states=True
      )
      assert np.array_equal(prob1, prob2)
      assert np.array_equal(prob1, prob5)
      assert np.array_equal(prob3, prob4)
      assert not np.array_equal(prob2, prob3)
  
file_path: deepchem/deepchem/rl/tests/test_a2c.py
imports: import unittest
  import pytest
  import numpy
  from flaky import flaky
  import deepchem
  from deepchem.models.optimizers import Adam, PolynomialDecay
label_desc: Use image node_id 2 for calling the TestA2C obj's underlying member method code with example usage: obj.test_recurrent_states() and returns: int, tf