ENH: Use Xarray for arrays with implicit structure

[mj023]

Is your feature request related to a problem?

Many arrays produced by PyLCM, like the value function array implicitly use a structure where each axis of the array corresponds to one dimension in the state- or state-action-space. The order of these axes is not fixed and for example depends on the order of the dimensions in the model formulation etc. It would be great to make this axis structure explicit.

Describe the solution you'd like

With Xarray it is possible generate multi-dimensional arrays with named axes, which still support many of the JAX Operations used in PyLCM. Therefore it might be possible to fully swap to Xarrays or, if we can not get them to work with vmap or other Jax functions, at least use them to add information to some the bigger arrays that users might want to interact with.

Benefits

• Simplifies user interactions with Vf_arr because it is clear which axis corresponds to which state space dimension (Logging, Debugging)
• Might make some code that keeps track of axis ordering obsolete
• Prevents axis mixups when developing