[BugFix] Sorting nested keys

[matteobettini]

Description

This PR solves issue #775

[codecov]

Codecov Report

Merging #787 (51ab9ab) into main (51ab9ab) will not change coverage.
The diff coverage is n/a.

❗ Current head 51ab9ab differs from pull request most recent head e714dd8. Consider uploading reports for the commit e714dd8 to get more accurate results

@@           Coverage Diff           @@
##             main     #787   +/-   ##
=======================================
  Coverage   88.78%   88.78%           
=======================================
  Files         123      123           
  Lines       21253    21253           
=======================================
  Hits        18869    18869           
  Misses       2384     2384           

Flag	Coverage Δ
habitat-gpu	24.79% <0.00%> (ø)
linux-brax	29.37% <0.00%> (ø)
linux-cpu	85.28% <0.00%> (ø)
linux-gpu	86.25% <0.00%> (ø)
linux-jumanji	30.15% <0.00%> (ø)
linux-outdeps-gpu	72.36% <0.00%> (ø)
linux-stable-cpu	85.14% <0.00%> (ø)
linux-stable-gpu	85.90% <0.00%> (ø)
linux_examples-gpu	42.68% <0.00%> (ø)
macos-cpu	85.04% <0.00%> (ø)
olddeps-gpu	76.12% <0.00%> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[vmoens]

LGTM, let's just make sure that we don't have keys that collide

[vmoens]

@matteobettini since the problem was occuring with Brax I updated the brax wrapper and tests to make sure that things were passing.
In general, it does not really make sense to execute brax in parallel envs (since it's vectorized), but it's a good way to spot bugs.

[matteobettini]

@matteobettini since the problem was occuring with Brax I updated the brax wrapper and tests to make sure that
things were passing.
In general, it does not really make sense to execute brax in parallel envs (since it's vectorized), but it's a good way to spot bugs.

@vmoens
I actually think it is important to use parallel envs with brax (and vmas).

This is because yes they perform batched simulation but parallelenv can furtherly distribute this to the various cpu cores.

This is also why rllib has the 2 parameters: num_workers and num_envs_per_worker