Cleanup imports

[aliberts]

What this does

~~Adds register mechanism to configs +~~ cleanup some imports

~~Related: #807~~

How it was tested

TODO

How to check it out

from dataclasses import dataclass

import draccus
from lerobot.common.policies import PreTrainedConfig


@dataclass
class MyConfig(PreTrainedConfig):
    n_layers: int = 5

    def observation_delta_indices(self): ...
    def action_delta_indices(self): ...
    def reward_delta_indices(self): ...
    def get_optimizer_preset(self): ...
    def get_scheduler_preset(self): ...
    def validate_features(self): ...


@dataclass
class MainConfig:
    policy: PreTrainedConfig


@draccus.wrap()
def main(cfg: MainConfig):
    print(cfg.policy)


if __name__ == "__main__":
    PreTrainedConfig.register_subclass("my_config", MyConfig)
    main()

[bsprenger]

Hi Simon,

Nice PR! :rocket: I saw that you tagged my PR, so I was wondering -- are you planning another way to use external configs with the training scripts, without needing to modify the imports? Or are the training scripts only supposed to be used with the configs already in the repo? Just wondering, so I know whether to update or close my PR.

Thanks for your insight!

[aliberts]

@bsprenger so there's a feature already built-in in draccus (PlugingRegistry) for registering external classes. This PR uses it instead of ChoiceRegistry.

I think we'll be able to combine yours on top of it since yours adds CLI parsing. But we should be able to register a new policies outside of LeRobot's training script (e.g. in the example above).

Let's put your PR in standby for now and we'll be able to merge it on top of this one ;)

[bsprenger]

Thanks for the reply, sounds good! 😀

[aliberts]

@bsprenger after further investigation, I think we don't need to use PluginRegistry as it's not doing exactly what we want. We can already use the register_subclass from the ChoiceRegistryBase class. We will however need to have lazy imports in the future as it's currently importing all policies, but that's going to be for another PR.

I'll finish this and then merge your PR.

[bsprenger]

Makes sense!

For the lazy imports, it might be nice if draccus itself supported discover_packages_path as a CLI arg (currently it has to be hardcoded in the plugin registry, I think?). It doesn't look easy to patch that in without some big changes to draccus though, which is why I proposed my PR as an intermediate step. But as you said, these are topics for a future MR. 😄

Really appreciate your comments!

[github-actions[bot]]

This PR has been automatically marked as stale because it has not had recent activity (1 year). It will be closed if no further activity occurs. Thank you for your contributions.

[github-actions[bot]]

This PR was closed because it has been stalled for 14 days with no activity. Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.