Use data classes to structure information and index gpudrive tensors by name. This also removes redundancy in the jax and torch implementations. Tested with puffer PPO.
Todo
[x] Ego state
[x] Partner obs
[x] LiDAR obs
[x] Log trajectory
[x] Road map global
[x] Road map local
also an important fix the PR this morning that collapses the reward tensor
Description
Use data classes to structure information and index
gpudrivetensors by name. This also removes redundancy in the jax and torch implementations. Tested with puffer PPO.Todo
[x] Ego state
[x] Partner obs
[x] LiDAR obs
[x] Log trajectory
[x] Road map global
[x] Road map local
also an important fix the PR this morning that collapses the reward tensor