code: def test_flatten_data_connector(self):
      ctx = ConnectorContext()
  
      c = FlattenDataAgentConnector(ctx)
  
      name, params = c.to_state()
      restored = get_connector(name, ctx, params)
      self.assertTrue(isinstance(restored, FlattenDataAgentConnector))
  
      sample_batch = {
          SampleBatch.NEXT_OBS: {
              "sensor1": [[1, 1], [2, 2]],
              "sensor2": 8.8,
          },
          SampleBatch.REWARDS: 5.8,
          SampleBatch.ACTIONS: [[1, 1], [2, 2]],
          SampleBatch.INFOS: {"random": "info"},
      }
  
      d = AgentConnectorDataType(
          0,
          1,
          # FlattenDataAgentConnector does NOT touch raw_dict,
          # so simply pass None here.
          AgentConnectorsOutput(None, sample_batch),
      )
  
      flattened = c([d])
      self.assertEqual(len(flattened), 1)
  
      batch = flattened[0].data.sample_batch
      self.assertTrue(
          (batch[SampleBatch.NEXT_OBS] == [1, 1, 2, 2, 8.8]).all()
      )
      self.assertEqual(batch[SampleBatch.REWARDS][0], 5.8)
      # Not flattened.
      self.assertEqual(len(batch[SampleBatch.ACTIONS]), 2)
      self.assertEqual(batch[SampleBatch.INFOS]["random"], "info")
  
imports: import gymnasium
  from gymnasium.spaces import Box
  import numpy
  import tree
  import unittest
  from ray.rllib.algorithms.ppo.ppo import PPOConfig
  from ray.rllib.connectors.agent.clip_reward import ClipRewardAgentConnector
  from ray.rllib.connectors.agent.lambdas import FlattenDataAgentConnector
  from ray.rllib.connectors.agent.obs_preproc import ObsPreprocessorConnector
  from ray.rllib.connectors.agent.pipeline import AgentConnectorPipeline
  from ray.rllib.connectors.agent.state_buffer import StateBufferConnector
  from ray.rllib.connectors.agent.view_requirement import ViewRequirementAgentConnector
  from ray.rllib.connectors.connector import ConnectorContext
  from ray.rllib.connectors.registry import get_connector
  from ray.rllib.policy.view_requirement import ViewRequirement
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils.test_utils import check
  from ray.rllib.utils.typing import ActionConnectorDataType, AgentConnectorDataType, AgentConnectorsOutput
  from ray.rllib.connectors.agent.mean_std_filter import MeanStdObservationFilterAgentConnector
label_desc: Use image node_id 4 for calling the TestAgentConnector obj's underlying member method code with example usage: obj.test_flatten_data_connector() without return types