mypy: Make Policies module MyPy‑compliant (#1719, #1720)

[akacmazz]

What this does

• Addresses Make LeRobot mypy compliant #1719; fixes Ensure the policy module passes MyPy type checks #1720. (🛠️ Tooling)
• Policies typing:
  • ACT: type‑safe activation helper; move batch tensors to the module device before normalization for consistent
    device semantics.
  • SAC: constructor accepts None to match tests (base class raises the expected error); forward returns a typed dict.
  • Config: PreTrainedConfig now respects LEROBOT_TEST_DEVICE for deterministic device selection in CI/local runs.
• Tooling (pre‑commit): enable a scoped MyPy run via the hook (uses pyproject.toml, installs types, ignores missing
  imports in the hook venv to avoid heavy deps). No global behavior change.

How it was tested

• Static checks:
  • pre-commit run -a → all hooks green (ruff, typos, bandit, mypy).
  • mypy → success on the policies scope configured in pyproject.toml.
• Tests (CPU to keep runs deterministic and lightweight):
  • SAC: PYTHONPATH=src pytest -q tests/policies -k sac
  • ACT back‑compat (uses small artifacts): PYTHONPATH=src pytest -q tests/policies -k 'act and
    backward_compatibility'
• Notes:
  • Set local HF caches to avoid permission issues and reduce flakiness.
  • On low‑memory machines, split ACT/SAC runs to avoid OOM.

How to checkout & try? (for the reviewer)

From repo root

export LEROBOT_TEST_DEVICE=cpu
mkdir -p .cache/hf_datasets .cache/hf_home
export HF_DATASETS_CACHE=$(pwd)/.cache/hf_datasets
export HF_HOME=$(pwd)/.cache/hf_home

Static checks

pre-commit run -a
mypy

SAC tests

PYTHONPATH=src pytest -q tests/policies -k sac

ACT back-compat tests

PYTHONPATH=src pytest -q tests/policies -k 'act and backward_compatibility'
If you’d like me to extend MyPy coverage to the next module (e.g., Datasets) in a follow‑up PR, I can prepare the scoped
config and initial fixes similarly.

cc @CadeRemi @AdilZouitine for review

[AdilZouitine]

--ignore-missing-import is true here instead of false, we should remove this flag right?

[akacmazz]

Thx for the review. :) Yes, we should remove --ignore-missing-imports.

[AdilZouitine]

Hey thank you for your contribution, I started the review for Act policy 😄
I will do the next policy later 🤗

[AdilZouitine]

cast should be always justified in the code, this is unsafe.
In this case it should be ok.

[AdilZouitine]

Ditto: we should always justify cast.
Question this cast is necessary?

[akacmazz]

Yes, i think so. this cast is necessary for MyPy compliance.

Without it, MyPy reports: Incompatible return value type (got "PreTrainedConfig", expected "T")

The issue is that while orig_config.__class__ and cls are the same type at runtime, MyPy can't statically infer that draccus.parse(orig_config.__class__, ...) will return the specific generic type T rather than the base PreTrainedConfig type.

This is a use of cast() to help the type checker understand a type relationship that's guaranteed at runtime but not statically provable.

[AdilZouitine]

Using object here feels too generic and makes the typing less helpful. A stricter approach would be:

Option A (inline type):

Sequence[dict[str, list[nn.Parameter] | float]]

Option B (preferred, clearer): define a dedicated TypedDict for optimizer kwargs and then use:

from typing import TypedDict, NotRequired
from collections.abc import Sequence
from torch import nn

class OptimizerKwargs(TypedDict, total=False):
    params: list[nn.Parameter]
    lr: float
    weight_decay: float
    momentum: float
    betas: tuple[float, float]
    eps: float

BatchKwargs = Sequence[OptimizerKwargs]

Option B has the advantage of being more self-documenting and makes it easier to see what fields are expected in the optimizer kwargs.

[akacmazz]

You're right.

I'll go with Option B (TypedDict) as you suggested. It makes it clear what the optimizer expects. I'll define the OptimizerKwargs TypedDict and use Sequence[OptimizerKwargs] as the return type.

Thanks for the detailed suggestion.

[AdilZouitine]

This addition is not related to mypy type?

[akacmazz]

Yes, you are right sorry for that. This device handling code is not related to MyPy typing. This was accidentally included from another change I was working on.

I'll remove this addition and keep this PR focused solely on type annotations and MyPy compliance.

Thanks for catching this.

[AdilZouitine]

**kwargs should be typed using typing.TypedDict and typing_extensions.Unpack.

[AdilZouitine]

ditto

[AdilZouitine]

Are cast necessary here?

[akacmazz]

These casts are unnecessary. "F.relu", "F.gelu", and "F.glu" are already properly typed as callables in PyTorch, so MyPy doesn't complain without the casts. I was being overly cautious here.

We can remove these casts and keep the function signatures clean:

def get_activation_fn(activation: str) -> Callable[..., Any]:
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu

Thanks for catching this unnecessary complexity. :)