[Feature] Added batch_lock attribute in EnvBase
#399
Conversation
vmoens
left a comment
There was a problem hiding this comment.
Perhaps we want to check that our envs have the attribute we want somewhere in the tests?
(like in test_env, create one gym and one model-based and assess whether they both have the correct batch_locked?)
Also we should test that this can't be set once that is implemented
|
I've added tests for gym. I'll add the same test for MB_env in the MBEnv PR #333 |
vmoens
left a comment
There was a problem hiding this comment.
I added a couple more comments, can you address these?
|
I think tests are failing because you have not merged main into your branch recently but i might be wrong |
vmoens
left a comment
There was a problem hiding this comment.
I'd like to know if you consider that TransformedEnv should make _batch_locked mutable? Can we also remove the commented lines?
|
Why would we make TransformedEnv being able to change the batch_locked. At the end of the day, base_env will still have that constraint |
vmoens
left a comment
There was a problem hiding this comment.
Sorry about this, still a couple of fix to be made before landing
batch_lock attribute in EnvBase
vmoens
left a comment
There was a problem hiding this comment.
almost there, a couple of minor things
|
@nicolas-dufour have a look at my changes, if you're ok with them we can merge |
Description
This PR solves the problem to distinguish if batch size of an env is flexible or not
Motivation and Context
The motivation is that stateful envs have different behavious than stateless envs with respect to batch_size. Where as stateful envs have a hard constraint on env.batch_size == tensordict.batch_size it's not the case of stateless envs. Therefore we add batch_locked attribute that allow or not batch flexibility in the env
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!