[NOMRG] Example: mass as Tensor

[matteobettini]

cc @LukasSchaefer

If you look at the files changed, here is an example of setting mass to a tensor (that is different for each env) in the navigation scenario

To run the scenario you can use

python vmas/examples/use_vmas_env.py

in this branch

It just happens to work out of the box for this case.

If you experience cases where it is not working out of the box let me know and we can fix them. Based on what you need we can tackle one attribute at a time to allow them to be union of float and tensor

[LukasSchaefer]

Thanks a lot @matteobettini ! This is really helpful to get started with the changes of vectorising some properties :) I'll likely continue further development in my fork but would be happy to look to merge some changes in should you want that/ be interested.

[matteobettini]

Thanks a lot @matteobettini ! This is really helpful to get started with the changes of vectorising some properties :) I'll likely continue further development in my fork but would be happy to look to merge some changes in should you want that/ be interested.

Absolutely! If there is any other properties that you feel would benefit from being Union[float,Tensor], feel free to either make a PR (which is super welcome) or let me know and i can see what I can do to make it possible

[matteobettini]

One tricky thing that might be worth being careful about is uninitended broadcasting:

For example if a property that was a float is now a tensor, its shape might have an impact on some operations.

# accell has shape [b,2]
# mass has shape [b, 1]
accell * mass
# the result has shape [b, 2]

# speed has shape [b]
# mass has shape [b, 1]
speed * mass
# the result has shape [b, b] -> we do not want this

So after changing a property to a tensor it might not be sufficient to just check it runs, but you might want to check that in the operations it is used in the shape of the output remains the same as if the property was a float

I checked this for mass and made sure that if you provide mass as a tensor of shape batch_dim, 1 it does not trigger unexpeceted bradcasting

[LukasSchaefer]

Yeah, great point on the broadcasting. I just ran into this. I agree with your reshaping suggestion, I think making sure everything is 2D in the form of batch_dim, 1 to avoid the exact issues you were describing is a clean way of addressing this issue. I'll test scenarios and see if I can identify any issues but will work on a PR after I have things sorted out :)

[matteobettini]

Yeah, great point on the broadcasting. I just ran into this. I agree with your reshaping suggestion, I think making sure everything is 2D in the form of batch_dim, 1 to avoid the exact issues you were describing is a clean way of addressing this issue. I'll test scenarios and see if I can identify any issues but will work on a PR after I have things sorted out :)

Yeah stuff in the backend might be bigger than 2D since we are batching agents as well in the vectorized constraint engine. But in general yes, we just have to make sure this issue does not arise for properties that we allow to be tensors