code: def test_weighted_generator(self):
      data0 = types.Transition(
          np.array([[1], [2], [3]]), (), _REWARD, (), ()
      )
      it0 = iter([data0])
  
      data1 = types.Transition(
          np.array([[4], [5], [6]]), (), _REWARD, (), ()
      )
      data2 = types.Transition(
          np.array([[7], [8], [9]]), (), _REWARD, (), ()
      )
      it1 = iter(
          [
              reverb.ReplaySample(
                  info=reverb.SampleInfo(
                      *[() for _ in reverb.SampleInfo.tf_dtypes()]
                  ),
                  data=data1,
              ),
              reverb.ReplaySample(
                  info=reverb.SampleInfo(
                      *[() for _ in reverb.SampleInfo.tf_dtypes()]
                  ),
                  data=data2,
              ),
          ]
      )
  
      weighted_it = builder._generate_samples_with_demonstrations(
          it0, it1, policy_to_expert_data_ratio=2, batch_size=3
      )
  
      np.testing.assert_array_equal(
          next(weighted_it).data.observation, np.array([[1], [4], [5]])
      )
      np.testing.assert_array_equal(
          next(weighted_it).data.observation, np.array([[7], [8], [2]])
      )
      self.assertRaises(StopIteration, lambda: next(weighted_it))
  
imports: from acme import types
  from acme.agents.jax.ail import builder
  import numpy
  import reverb
  from absl.testing import absltest
label_desc: Use image node_id 1 for calling the BuilderTest obj's underlying member method code with example usage: obj.test_weighted_generator() without return types