code: def shuffled_batches(self, batch_size):
          
          if batch_size >= self.num_envs * self.num_steps:
              yield self
          else:
              rollouts_in_batch = batch_size // self.num_steps
  
              batch_splits = math_util.divide_ceiling(self.num_envs, rollouts_in_batch)
  
              indices = list(range(self.num_envs))
              np.random.shuffle(indices)
  
              for sub_indices in np.array_split(indices, batch_splits):
                  yield Trajectories(
                      num_steps=self.num_steps,
                      num_envs=len(sub_indices),
                      # Dont use it in batches for a moment, can be uncommented later if needed
                      # environment_information=[x[sub_indices.tolist()] for x in self.environment_information],
                      environment_information=None,
                      transition_tensors={k: x[:, sub_indices] for k, x in self.transition_tensors.items()},
                      rollout_tensors={k: x[sub_indices] for k, x in self.rollout_tensors.items()},
                      # extra_data does not go into batches
                  )
docs: Generate randomized batches of data - only sample whole trajectories