Fix: Keep reward shape and dtype the same when resetting and stepping

[RuanJohn]

What

Due to default behaviour in Jumanji the reward was set to a single float value when the environment was reset, but when stepping num_agent int rewards are returned. This PR fixes this by passing in num_agents as the shape argument to the restart, termination and transition methods in Jumanji.

Extra

• Added a new test that checks that all leaves in the timestep pytree have the same shapes and data types when resetting and stepping the environment.
• Since discounts will now also have shape (num_agents, ) a test checking the discount shapes was also updated.
• Data types of the rewards are also explicitly cast to floats to ensure consistency between stepping and resetting the environment.

[arnupretorius]

Thanks @RuanJohn 👍