Skip to content

Conversation

@akacmazz
Copy link

@akacmazz akacmazz commented Oct 23, 2025

What this does

This PR makes the lerobot.policies module 100% MyPy compliant by enabling strict type checking and adding comprehensive type annotations to all policy implementations.

Addresses #1719 ; fixes #1720 . (🔧 Enhancement)

Changes:

MyPy Configuration (pyproject.toml):

  • Enabled strict MyPy checking for all 8 policy modeling modules
  • Configured gradual typing approach with module-specific overrides
  • Added disallow_subclassing_any = false for nn.Module compatibility

Base Classes (src/lerobot/policies/pretrained.py):

  • Fixed ClassVar annotations for config_class and name attributes
  • Updated get_optim_params() return type to dict | list[dict] for flexibility

Policy Implementations (50+ type hints added across 8 files):

  • ACT (11 type hints): Return type annotations for __init__, reset, helper methods
  • Diffusion (13 type hints): Fixed **kwargs typing, added parameter types
  • TDMPC (8 type hints): Return types for __init__, reset, update methods
  • SAC (8 type hints): Return types for main class and helper classes
  • Pi0 (5 type hints): Return types for policy and helper classes
  • Pi05 (5 type hints): Return types for policy and helper classes
  • SmolVLA (3 type hints): Return types for SmolVLAPolicy
  • VQBeT (6 type hints): Return types for policy and helper classes

Type Checking Fixes:

  • Fixed None type narrowing with assertions
  • Added 12 strategic type: ignore comments with TODO(Ensure the policy module passes MyPy type checks #1720) explanations for genuine edge cases:
    • **kwargs unpacking limitations (6 instances in Diffusion)
    • Literal type mismatches with config strings (2 instances in Pi0/Pi05)
    • Intentional signature overrides (1 instance in SAC)
    • torch.compile method assignment (1 instance in Pi0)
    • getattr dynamic type resolution (1 instance in SAC)
    • tuple unpacking with * operator (1 instance in SmolVLA)

How it was tested

1. MyPy Type Checking (Primary Test):

mypy src/lerobot/policies --config-file=pyproject.toml
# Success: no issues found in 49 source files ✅

Result: 131 MyPy errors → 0 errors 

2. Python Syntax Validation:
python3 -m compileall src/lerobot/policies/ -q
# No compilation errors ✅

3. Type Ignore Documentation:
grep -B1 "type: ignore" src/lerobot/policies/*/modeling_*.py | grep "TODO"
# All 12 instances have TODO(#1720) explanations ✅

4. Compatibility:
- Successfully rebased onto latest main (includes GR00T model from PR #1452)
- All new code passes MyPy checks including newly added policies
- Changes are type-annotation-only; no functional code modifications

How to checkout & try? (for the reviewer)

1. Verify MyPy passes on the policies module:
mypy src/lerobot/policies --config-file=pyproject.toml
Expected output: Success: no issues found in 49 source files

2. Check specific policy files:
mypy src/lerobot/policies/act/modeling_act.py --config-file=pyproject.toml
mypy src/lerobot/policies/diffusion/modeling_diffusion.py --config-file=pyproject.toml
mypy src/lerobot/policies/sac/modeling_sac.py --config-file=pyproject.toml

3. Verify all type: ignore comments have explanations:
grep -B1 "type: ignore" src/lerobot/policies/*/modeling_*.py | grep "TODO"
All 12 instances should show TODO(#1720) with explanations.

4. Review key changes:
- src/lerobot/policies/pretrained.py:49-50 - ClassVar annotations
- src/lerobot/policies/pretrained.py:160 - get_optim_params return type
- pyproject.toml:346-369 - MyPy configuration for policies module

5. Verify Python compilation:
python3 -m compileall src/lerobot/policies/ -q
This commit enables strict MyPy type checking for the policies module
and adds comprehensive type annotations to ACT, Diffusion, and TDMPC
policies.

Changes:
- pyproject.toml: Enable strict MyPy for policies module
- ACT policy (11 type hints): Added return type annotations to __init__,
  reset, and helper methods; fixed Callable return type
- Diffusion policy (13 type hints): Added return type annotations,
  fixed **kwargs typing, added parameter types to helper classes
- TDMPC policy (8 type hints): Added return type annotations to __init__,
  reset, update, and internal methods

Related to issue huggingface#1720
Added 8 return type annotations to SAC policy:
- Main SACPolicy class: __init__, reset
- Update methods: update_target_networks, update_temperature
- Helper classes: MLP, CriticHead, CriticEnsemble, DiscreteCritic, Policy

Related to issue huggingface#1720
Added 10 return type annotations total (5 each for Pi0 and Pi05):
- Pi0Policy and PI05Policy classes: __init__, reset
- PI0Pytorch and PI05Pytorch: __init__
- PaliGemmaWithExpertModel: __init__ with typed parameters
- GemmaConfig: __init__ with all int parameters

Related to issue huggingface#1720
Added 9 return type annotations total:
- SmolVLA (3 type hints): SmolVLAPolicy __init__ and reset, VLAFlowMatching __init__
- VQBeT (6 type hints): VQBeTPolicy __init__ and reset, SpatialSoftmax, VQBeTModel, VQBeTHead, VQBeTRgbEncoder

Related to issue huggingface#1720
Fixed 113 MyPy errors (131 → 18):
- Added type annotations for queues and action queues across all policies
- Fixed None type narrowing with assertions in ACT, Diffusion, SAC, VQBeT
- Fixed return type annotations (get_optim_params)
- Added 'from typing import Any' imports where needed
- Enabled MyPy subclassing of Any types in pyproject.toml

Remaining 18 errors are complex edge cases in SAC, Pi0/Pi05, SmolVLA.

Related to issue huggingface#1720
Fixed all remaining 18 MyPy errors:
- Added type annotations for inputs_embeds and adarms_cond in Pi0/Pi05
- Fixed variable redefinition errors in SAC by removing duplicate type annotations
- Fixed SmolVLA forward return type (dict → tuple[Tensor, dict])
- Added type: ignore for intentional cases (torch.compile, getattr, resize_with_pad)
- Added missing 'from typing import Any' imports

RESULT: MyPy now passes with 0 errors on all 8 policy files!
Starting point: 131 errors → Final: 0 errors (100% reduction)

Related to issue huggingface#1720
Fixed all remaining 18 MyPy errors:
- Added type annotations for inputs_embeds and adarms_cond in Pi0/Pi05
- Fixed variable redefinition errors in SAC by removing duplicate type annotations
- Fixed SmolVLA forward return type (dict → tuple[Tensor, dict])
- Added type: ignore for intentional cases (torch.compile, getattr, resize_with_pad)
- Added missing 'from typing import Any' imports

RESULT: MyPy now passes with 0 errors on all 8 policy files!
Starting point: 131 errors → Final: 0 errors (100% reduction)

Related to issue huggingface#1720
Added TODO(huggingface#1720) comments with clear explanations for all 12 type: ignore instances as required by Issue huggingface#1720.

Changes:
- Diffusion (6 instances): MyPy cannot infer **kwargs types when unpacking dict
- Pi0 (2 instances): config.dtype Literal mismatch, torch.compile method assignment
- Pi05 (1 instance): config.dtype Literal mismatch
- SAC (2 instances): Intentional forward() override, getattr runtime type resolution
- SmolVLA (1 instance): MyPy cannot infer tuple unpacking with * operator

All type: ignore comments now include:
- TODO(huggingface#1720) reference linking to the issue
- Clear explanation of why type: ignore is necessary

MyPy verification: Success - no issues found in 49 source files
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

1 participant