code/__init__.py

"""Recursive latent planner over the frozen LeWM PushT world model.

Phase 2 of the project. The recursion replaces the plan-token refiner of
``lejepa_control.controller.IterativeController``; it does not replace the
objective around it, so the arrival-and-hold term, the late path weighting, the
GMM support model and the dataset loader are imported from phase 1 unchanged.

The phase-1 controller is left untouched — it is the baseline this planner is
measured against.
"""

from lejepa_control_2.losses import (
    DistanceScale,
    anchor_loss,
    cycle_weights,
    deep_supervision_loss,
    planner_loss,
    saturation_loss,
)
from lejepa_control_2.planner import (
    ActionEmbedding,
    FTheta,
    GTheta,
    RecursivePlanner,
    RMSNorm,
    consequence_features,
    m_step,
)

__all__ = [
    'ActionEmbedding',
    'DistanceScale',
    'FTheta',
    'GTheta',
    'RMSNorm',
    'RecursivePlanner',
    'anchor_loss',
    'consequence_features',
    'cycle_weights',
    'deep_supervision_loss',
    'm_step',
    'planner_loss',
    'saturation_loss',
]