complaints: 1. When generating u_traj for CEM, we need a leaf function which takes the mean and covariance. Right now, this isn't possible with generate_batch, which means we need to use tree_map and cast the result as a batch, which is a bit disappointing. 2. When adding a traj to a trajbatch, we need to cast to batch first. That's also annoying. 3. step, output, and cost need lots of case handling for the traj/batch stuf Goal is to adjust generate_batch or add another function which can handle such initialization cases.