code: def flatten_data(data: AgentConnectorsOutput):
      assert isinstance(
          data, AgentConnectorsOutput
      ), "Single agent data must be of type AgentConnectorsOutput"
  
      raw_dict = data.raw_dict
      sample_batch = data.sample_batch
  
      flattened = {}
      for k, v in sample_batch.items():
          if k in [
              SampleBatch.INFOS,
              SampleBatch.ACTIONS,
          ] or k.startswith("state_out_"):
              # Do not flatten infos, actions, and state_out_ columns.
              flattened[k] = v
              continue
          if v is None:
              # Keep the same column shape.
              flattened[k] = None
              continue
          flattened[k] = np.array(tree.flatten(v))
      flattened = SampleBatch(flattened, is_training=False)
  
      return AgentConnectorsOutput(raw_dict, flattened)
  
file_path: ray/rllib/connectors/agent/lambdas.py
imports: from typing import Any, Callable, Type
  import numpy
  import tree
  from ray.rllib.connectors.connector import AgentConnector, ConnectorContext
  from ray.rllib.connectors.registry import register_connector
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils.typing import AgentConnectorDataType, AgentConnectorsOutput
  from ray.util.annotations import PublicAPI
label_desc: Use image node_id 2 for calling a global function with example usage: flatten_data(data) and returns: AgentConnectorsOutput