docstr: Tests, whether Policy/Model return correct LSTM ViewRequirements.
docstr_tok: ["Tests",",","whether","Policy\/Model","return","correct","LSTM","ViewRequirements","."]
code: def test_traj_view_lstm_prev_actions_and_rewards(self):
      """Tests, whether Policy/Model return correct LSTM ViewRequirements."""
  
      # This test only works on ModelV2 stack. So, we need to disable the RLModule
      # and Learner API.
      config = (
          ppo.PPOConfig()
          .experimental(_enable_new_api_stack=False)
          .environment("CartPole-v1")
          # Activate LSTM + prev-action + rewards.
          .training(
              model={
                  "use_lstm": True,
                  "lstm_use_prev_action": True,
                  "lstm_use_prev_reward": True,
              },
          )
          .rollouts(create_env_on_local_worker=True)
      )
  
      for _ in framework_iterator(config):
          algo = config.build()
          policy = algo.get_policy()
          view_req_model = policy.model.view_requirements
          view_req_policy = policy.view_requirements
          # 7=obs, prev-a + r, 2x state-in, 2x state-out.
          assert len(view_req_model) == 7, view_req_model
          assert len(view_req_policy) == 23, (
              len(view_req_policy),
              view_req_policy,
          )
          for key in [
              SampleBatch.OBS,
              SampleBatch.ACTIONS,
              SampleBatch.REWARDS,
              SampleBatch.TERMINATEDS,
              SampleBatch.TRUNCATEDS,
              SampleBatch.NEXT_OBS,
              SampleBatch.VF_PREDS,
              SampleBatch.PREV_ACTIONS,
              SampleBatch.PREV_REWARDS,
              "advantages",
              "value_targets",
              SampleBatch.ACTION_DIST_INPUTS,
              SampleBatch.ACTION_LOGP,
          ]:
              assert key in view_req_policy
  
              if key == SampleBatch.PREV_ACTIONS:
                  assert (
                      view_req_policy[key].data_col
                      == SampleBatch.ACTIONS
                  )
                  assert view_req_policy[key].shift == -1
              elif key == SampleBatch.PREV_REWARDS:
                  assert (
                      view_req_policy[key].data_col
                      == SampleBatch.REWARDS
                  )
                  assert view_req_policy[key].shift == -1
              elif key not in [
                  SampleBatch.NEXT_OBS,
                  SampleBatch.PREV_ACTIONS,
                  SampleBatch.PREV_REWARDS,
              ]:
                  assert view_req_policy[key].data_col is None
              else:
                  assert (
                      view_req_policy[key].data_col == SampleBatch.OBS
                  )
                  assert view_req_policy[key].shift == 1
  
          rollout_worker = algo.workers.local_worker()
          sample_batch = rollout_worker.sample()
          sample_batch = convert_ma_batch_to_sample_batch(sample_batch)
  
          # Rollout fragment length should be auto-computed to 2000:
          # 2 workers, 1 env per worker, train batch size=4000 -> 2000 per worker.
          self.assertEqual(
              sample_batch.count, 2000, "ppo rollout count != 2000"
          )
          self.assertEqual(
              sum(sample_batch["seq_lens"]), sample_batch.count
          )
          self.assertEqual(
              len(sample_batch["seq_lens"]),
              sample_batch["state_in_0"].shape[0],
          )
  
          # check if non-zero state_ins are pointing to the correct state_outs
          seq_counters = np.cumsum(sample_batch["seq_lens"])
          for i in range(sample_batch["state_in_0"].shape[0]):
              state_in = sample_batch["state_in_0"][i]
              if np.any(state_in != 0):
                  # non-zero state-in should be one of th state_outs.
                  state_out_ind = seq_counters[i - 1] - 1
                  check(
                      sample_batch["state_out_0"][state_out_ind],
                      state_in,
                  )
          algo.stop()
  
imports: import gymnasium
  from gymnasium.spaces import Box, Discrete
  import numpy
  import unittest
  import ray
  from ray.rllib.algorithms.callbacks import DefaultCallbacks
  import ray.rllib.algorithms.dqn
  import ray.rllib.algorithms.ppo
  from ray.rllib.examples.env.debug_counter_env import MultiAgentDebugCounterEnv
  from ray.rllib.examples.env.multi_agent import MultiAgentPendulum
  from ray.rllib.evaluation.rollout_worker import RolloutWorker
  from ray.rllib.examples.policy.episode_env_aware_policy import EpisodeEnvAwareAttentionPolicy, StatefulRandomPolicy
  from ray.rllib.models.tf.attention_net import GTrXLNet
  from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch, convert_ma_batch_to_sample_batch
  from ray.rllib.policy.view_requirement import ViewRequirement
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.test_utils import framework_iterator, check
label_desc: Use image node_id 4 for calling the TestTrajectoryViewAPI obj's underlying member method code with example usage: obj.test_traj_view_lstm_prev_actions_and_rewards() without return types