Improve splitting of trajectory datasets.

[cisprague]

We need to be able to succinctly split a trajectory dataset into train/test/val/etc by trajectories or possibly other axes.

E.g., assume we have a (n_trajs, n_steps, n_dims) dataset of trajectories (assumes padded).

• We could split into train/val along the n_trajs axis.

• We could reshape the dataset to be of adjacent states of shape(n_trajs*n_steps // 2, 2*n_dims) and then split along the first axis.

• Also worth looking at how the well handles different splits.

[sgreenbury]

An approach to this is included in #840 with the data module approach from the-well. This involves constructing a dataset for each of train, valid, test, rollout valid (free running for whole traj), rollout test (free running for whole traj).

This can be init with both data from memory or from disk assuming the datasets are present in specific path structure. It might be useful to add a helper method on AutoEmulateDataset that directly does the split during initialization and returns three datasets instead to provide additional flexibility so that the simulator just needs to be run once or write to one output path.