WIP : Add possibility for multitask training on ACT policy

[marinabar]

What this does

Implements the possibility to train on multiple datasets and creates a dataset_index token that conditions a FiLM layer applied to the inputs. That layer is applied to the output of the ResNet image feature extractor.

This layer is added if dataset_index is specified in the input_shapes.

This also adds an input field for dataset index upon evaluation in eval.py and control_robot.py if the yaml configuration is for multitask.

TODO?

• [ ] maybe : add evaluation on one or all datasets during training

How it was tested

Trained a multitask policy both on the insertion and transfer cube Aloha tasks, with the following hyperparameters :

• BS 16
• 4 encoder layers
• 1e-4 learning

https://wandb.ai/marinabar/lerobot/runs/anc0wly9?nw=nwuserm1bn

python lerobot/scripts/train.py \
 hydra.job.name=aloha_film_on_feature_multitask \
 hydra.run.dir=/fsx/marina_barannikov/outputs/multitask_08/aloha_film_on_feature_multitask \
 policy=act_multitask \
 env=aloha \
 wandb.enable=true

Performance :

WIP

How to checkout & try? (for the reviewer)

Provide a simple way for the reviewer to try out your changes.

Examples:

 python lerobot/scripts/train.py \
 hydra.job.name=film_multitask_aloha \
 policy=act \
 env=aloha \
 +policy.input_shapes.dataset_index=[1] \
 training.offline_steps=1000