[Feature] Brax envs #722
[Feature] Brax envs #722
Conversation
facebook-github-bot
added
the
CLA Signed
Codecov Report
@@ Coverage Diff @@
## main #722 +/- ##
==========================================
+ Coverage 88.72% 88.79% +0.06%
==========================================
Files 120 122 +2
Lines 20241 20492 +251
==========================================
+ Hits 17958 18195 +237
- Misses 2283 2297 +14
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
left a comment
vmoens
left a comment
There was a problem hiding this comment.
Let's set up the test pipeline, I'll give it a second look once it's done
| >>> env = brax.envs.get_environment("ant") | ||
| >>> env = BraxWrapper(env) | ||
| >>> td = env.rand_step() | ||
| >>> print(td) |
There was a problem hiding this comment.
we should put what the print output is here below
torchrl/envs/libs/brax.py
Outdated
| NdUnboundedContinuousTensorSpec, | ||
| ) | ||
| from torchrl.envs.common import _EnvWrapper | ||
| from torchrl.envs.libs.jumanji import _torchrl_data_to_spec_transform |
There was a problem hiding this comment.
we should definitely not import stuff from lib to lib.
Here I would put _torchrl_data_to_spec_transform in a utils file, jax_utils (and rename it to a less generic name).
|
The tests are passing! Just a few edits and we're good to go! |
vmoens
left a comment
vmoens
left a comment
There was a problem hiding this comment.
LGTM
Let's wait for the tests to pass and merge this!
3 participants
Description
Draft PR of Brax envs