code: def _forward_train(self, batch: NestedDict):
      output = {}
  
      # Shared encoder.
      encoder_outs = self.encoder(batch)
      if STATE_OUT in encoder_outs:
          output[STATE_OUT] = encoder_outs[STATE_OUT]
  
      # Value head.
      vf_out = self.vf(encoder_outs[ENCODER_OUT][CRITIC])
      # Squeeze out last dim (value function node).
      output[SampleBatch.VF_PREDS] = tf.squeeze(vf_out, axis=-1)
  
      # Policy head.
      action_logits = self.pi(encoder_outs[ENCODER_OUT][ACTOR])
      output[SampleBatch.ACTION_DIST_INPUTS] = action_logits
  
      return output
  
file_path: ray/python/ray/rllib/algorithms/ppo/tf/ppo_tf_rl_module.py
imports: from typing import Any, Dict
  from ray.rllib.algorithms.ppo.ppo_rl_module import PPORLModule
  from ray.rllib.core.models.base import ACTOR, CRITIC
  from ray.rllib.core.models.tf.encoder import ENCODER_OUT, STATE_OUT
  from ray.rllib.core.rl_module.rl_module import RLModule
  from ray.rllib.core.rl_module.tf.tf_rl_module import TfRLModule
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.nested_dict import NestedDict
label_desc: Use image node_id 3 for calling the PPOTfRLModule obj's underlying member method code with example usage: obj._forward_train(batch) and returns: output