Bug fix: missing attention mask in VAE encoder in ACT policy

[thomwolf]

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.

• Added a parametrized test ("aloha", "act_1000_actions", []), in tests/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:

• editing the saving format for the backward policy saved safetensors so that they include the overriden arguments in the naming and we could then use an override argument instead of a new policy (currently doing that overide the previous aloha/act test
• editing the previous aloha/act test to test the attention masking itself ( having both tested in the same run) by using the 1000 predicted steps instead of the 10 step previously tested.

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:

CUDA_VISIBLE_DEVICES="" pytest ./tests/test_policies.py::test_backward_compatibility\[aloha-act-extra_overrides2\]

---

This change is 

[alexander-soare]

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"].

[alexander-soare]

As you mentioned in your PR message, I think we can avoid adding a whole config file by doing

Suggested change

	("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?

[thomwolf]

yes in this case we still need a act_1000_actions config file...

[thomwolf]

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.

[alexander-soare]

Oh sorry, I meant also to change to using just the act.json file. Also happy to hear what @aliberts says.

[thomwolf]

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

[alexander-soare]

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

[aliberts]

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.

[alexander-soare]

Oh sorry, I meant also to change to using just the act.json file. Also happy to hear what @aliberts says.

[thomwolf]

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 (tests/test_policies.py::test_backward_compatibility[xarm-tdmpc-extra_overrides0])

Should be ready to merge whenever you think it make sense