-
Notifications
You must be signed in to change notification settings - Fork 641
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Bug fix: missing attention mask in VAE encoder in ACT policy #279
Bug fix: missing attention mask in VAE encoder in ACT policy #279
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks a lot for this @thomwolf. I've been waiting for this one as I want to test something out.
Regarding the test, do you know if the huge chunk size slows things down significantly? If so, maybe we can also make the model a lot smaller. See my comment about passing overrides via Python rather than adding a whole new config. Then we'd need to do something like ["policy.dim_model=64", "policy.dim_feedforward=64", "policy.n_encoder_layers=2", "policy.n_vae_encoder_layers=2"]
.
# "diffusion", | ||
# ["policy.n_action_steps=8", "policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"], | ||
# ), | ||
("aloha", "act_1000_actions", []), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
As you mentioned in your PR message, I think we can avoid adding a whole config file by doing
("aloha", "act_1000_actions", []), | |
("aloha", "act_1000_actions", ["policy.n_action_steps=1000", "policy.chunk_size=1000"]), |
(and in the test)
Feels like a no-brainer to me but maybe I'm missing something?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
yes in this case we still need a act_1000_actions
config file...
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe @aliberts you have an opinion on this question around tests and config file.
See my How it was tested
section in this PR description.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Oh sorry, I meant also to change to using just the act.json file. Also happy to hear what @aliberts says.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
oh yes this also doesn't work because the safetensor files overide themselves (cf my comment in the PR description. maybe a bit short) unless we either:
- decide we test always in this 1000 action setup for the other ACT policy tests as well
- decide to update the way we name safetensor files
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Got it, thanks. My stance on this is it's better to add test artifacts than to add config files to the source code. That can even mean just moving the yaml file to the artifacts directory. Also happy with both of your proposals. Waiting on @aliberts
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
unless we either:
- decide we test always in this 1000 action setup for the other ACT policy tests as well
- decide to update the way we name safetensor files
IMO both are fine (the naming trick you went with is okay, although it adds artifacts they're relatively small). I was concerned with speed too but from what I'm seeing in CI tests it doesn't increase tests duration (cf this branch vs main)
Long-term though, I think it'll be much better to have more fine-grained tests and test individual components of the policies and do away with these artifacts, similar to what's done in transformers. My motivations for this are:
- Whenever this test breaks, it doesn't tell us which part broke it.
- Having this many tests artifacts as we do right now is really not great for cloning the repo.
- It's kernel-dependent.
For now it's okay because this granularity allows us to iterate faster while still having guardrails. Happy to hear your opinions on this as well.
# "diffusion", | ||
# ["policy.n_action_steps=8", "policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"], | ||
# ), | ||
("aloha", "act_1000_actions", []), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Oh sorry, I meant also to change to using just the act.json file. Also happy to hear what @aliberts says.
Ok updated so that test fixtures are clean now. There is one failing test but it's independent of this PR and the same which is also failing on main ( Should be ready to merge whenever you think it make sense |
What this does
This pull request fix the missing masking in the VAE encoder in ACT modeling.
When we train, we sample various length of episodes so we have a mask and padding tokens when a sample is less long than the sequence of action to predict. We need to handle this padding for correct VAE encoding during training.
How it was tested
Tested the changes by adding an additional
test_backward_compatibility
configuration where the number of predicted action steps is longer than the length of the episodes in the dataset, hence padding is used for every sample.("aloha", "act_1000_actions", []),
intests/test_policies.py
and an assiciated configuration.Note that there could be other way to test this which would not involve adding a new configuation env in the lib, in particular:
I'm quite agnostic to the way we decide to organize the tests so happy to change one way or another.
How to checkout & try? (for the reviewer)
Examples:
This change is