code: def test_forward_train(self):
      # TODO: Add FrozenLake-v1 to cover LSTM case.
      frameworks = ["tf2", "torch"]
      env_names = ["CartPole-v1", "Pendulum-v1", "ALE/Breakout-v5"]
      lstm = [False, True]
      config_combinations = [frameworks, env_names, lstm]
      for config in itertools.product(*config_combinations):
          fw, env_name, lstm = config
          print(f"[FW={fw} | [ENV={env_name}] | LSTM={lstm}")
          env = gym.make(env_name)
  
          preprocessor_cls = get_preprocessor(env.observation_space)
          preprocessor = preprocessor_cls(env.observation_space)
  
          module = _get_ppo_module(
              framework=fw,
              env=env,
              lstm=lstm,
              observation_space=preprocessor.observation_space,
          )
  
          # collect a batch of data
          batches = []
          obs, _ = env.reset()
          obs = preprocessor.transform(obs)
          tstep = 0
  
          if lstm:
              state_in = module.get_initial_state()
              if fw == "torch":
                  state_in = tree.map_structure(
                      lambda x: x[None],
                      convert_to_torch_tensor(state_in),
                  )
              else:
                  state_in = tree.map_structure(
                      lambda x: tf.convert_to_tensor(x)[None], state_in
                  )
              initial_state = state_in
  
          while tstep < 10:
              input_batch = _get_input_batch_from_obs(
                  fw, obs, lstm=lstm
              )
              if lstm:
                  input_batch[STATE_IN] = state_in
  
              fwd_out = module.forward_exploration(input_batch)
              action_dist_cls = module.get_exploration_action_dist_cls()
              action_dist = action_dist_cls.from_logits(
                  fwd_out[SampleBatch.ACTION_DIST_INPUTS]
              )
              _action = action_dist.sample()
              action = convert_to_numpy(_action[0])
              action_logp = convert_to_numpy(
                  action_dist.logp(_action)[0]
              )
              if lstm:
                  # Since this is inference, fwd out should only contain one action
                  assert len(action) == 1
                  action = action[0]
              new_obs, reward, terminated, truncated, _ = env.step(
                  action
              )
              new_obs = preprocessor.transform(new_obs)
              output_batch = {
                  SampleBatch.OBS: obs,
                  SampleBatch.NEXT_OBS: new_obs,
                  SampleBatch.ACTIONS: action,
                  SampleBatch.ACTION_LOGP: action_logp,
                  SampleBatch.REWARDS: np.array(reward),
                  SampleBatch.TERMINATEDS: np.array(terminated),
                  SampleBatch.TRUNCATEDS: np.array(truncated),
                  STATE_IN: None,
              }
  
              if lstm:
                  assert STATE_OUT in fwd_out
                  state_in = fwd_out[STATE_OUT]
              batches.append(output_batch)
              obs = new_obs
              tstep += 1
  
          # convert the list of dicts to dict of lists
          batch = tree.map_structure(lambda *x: np.array(x), *batches)
          # convert dict of lists to dict of tensors
          if fw == "torch":
              fwd_in = {
                  k: convert_to_torch_tensor(np.array(v))
                  for k, v in batch.items()
              }
              if lstm:
                  fwd_in[STATE_IN] = initial_state
                  # If we test lstm, the collected timesteps make up only one batch
                  fwd_in = {
                      k: torch.unsqueeze(v, 0) if k != STATE_IN else v
                      for k, v in fwd_in.items()
                  }
  
              # forward train
              # before training make sure module is on the right device
              # and in training mode
              module.to("cpu")
              module.train()
              fwd_out = module.forward_train(fwd_in)
              loss = dummy_torch_ppo_loss(module, fwd_in, fwd_out)
              loss.backward()
  
              # check that all neural net parameters have gradients
              for param in module.parameters():
                  self.assertIsNotNone(param.grad)
          else:
              fwd_in = tree.map_structure(
                  lambda x: tf.convert_to_tensor(x, dtype=tf.float32),
                  batch,
              )
              if lstm:
                  fwd_in[STATE_IN] = initial_state
                  # If we test lstm, the collected timesteps make up only one batch
                  fwd_in = {
                      k: tf.expand_dims(v, 0) if k != STATE_IN else v
                      for k, v in fwd_in.items()
                  }
  
              with tf.GradientTape() as tape:
                  fwd_out = module.forward_train(fwd_in)
                  loss = dummy_tf_ppo_loss(module, fwd_in, fwd_out)
              grads = tape.gradient(loss, module.trainable_variables)
              for grad in grads:
                  self.assertIsNotNone(grad)
  
imports: import itertools
  import unittest
  import gymnasium
  import numpy
  import tree
  import ray
  from ray.rllib import SampleBatch
  from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
  from ray.rllib.algorithms.ppo.tf.ppo_tf_rl_module import PPOTfRLModule
  from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule
  from ray.rllib.core.models.base import STATE_IN, STATE_OUT
  from ray.rllib.core.rl_module.rl_module import RLModuleConfig
  from ray.rllib.models.preprocessors import get_preprocessor
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.torch_utils import convert_to_torch_tensor
label_desc: Use image node_id 4 for calling the TestPPO obj's underlying member method code with example usage: obj.test_forward_train() without return types