Fix padding shape in split_and_pad_trajectories to support arbitrary additional dimensions

[bikcrum]

Description

This PR fixes an issue in split_and_pad_trajectories(tensor, dones) where creating a full-length trajectory incorrectly assumes the input tensor is 3D ([time, number of envs, feature_dim]). This assumption fails for higher-dimensional tensors, such as images ([time, number of envs, channel, width, height]), where there are multiple additional dimensions. The previous approach, which used only the last dimension (shape[-1]), caused shape mismatches when adding a full-length zero trajectory.

Key Changes

Instead of assuming a single feature dimension (shape[-1]), the full-length trajectory is now created using all feature dimensions beyond the second dimension (number of environments), represented by *shape[2:].
This ensures compatibility with tensors of arbitrary additional dimensions while preserving the intended functionality.

For reference, the following line is changed:

rsl_rl/rsl_rl/utils/utils.py

Line 66 in f80d475

trajectories = trajectories + (torch.zeros(tensor.shape[0], tensor.shape[-1], device=tensor.device),)

[ClemensSchwarke]

Hey bikcrum, thanks a lot for opening a PR to fix this! The issue is already fixed on our internal development branch and will be resolved in the next release.
Thanks again :)