code: def setUp(self):
      super().setUp()
      self.trajectories = jnp.array(
          [  # [T, B]
              [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
              [0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
          ]
      ).transpose()
      self.start_indices = jnp.array(
          [[0, 1, 4], [1, 2, 5]]  # [B, num_starts]
      )
      self.invalid_start_indices = jnp.array(
          [  # [B, num_starts]
              [0, 1, 25],  # 25 is out of bound
              [1, 2, 5],
          ]
      )
  
file_path: rlax/rlax/_src/model_learning_test.py
imports: from absl.testing import absltest
  import jax.numpy
  import numpy
  from rlax._src import model_learning
label_desc: Use image node_id 1 for calling the ModelLearningTest obj's underlying member method code with example usage: obj.setUp() without return types