<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def shuffled_batches(self, batch_size):
 """ Generate randomized batches of data - only sample whole trajectories """

        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
                )