rewards/__init__.py

"""
Visual Memory reward system — two reward types for evaluation and RL training.

Custom Checks (episode-level):
    from rewards.checks import VisualMemoryChecker

OpenEnV Transforms (per-step):
    from rewards.transforms import VisualMemoryStepTransform

Base utilities:
    from rewards.base import Scenario, EpisodeLog, RewardCalculator, OpenEnvRewardCalculator
"""

from .checks import VisualMemoryChecker
from .transforms import VisualMemoryStepTransform
from .base import (
    Scenario,
    EpisodeLog,
    StepLog,
    RewardBreakdown,
    RewardCalculator,
    StepRewardTransform,
    OpenEnvRewardCalculator,
)

__all__ = [
    "VisualMemoryChecker",
    "VisualMemoryStepTransform",
    "Scenario",
    "EpisodeLog",
    "StepLog",
    "RewardBreakdown",
    "RewardCalculator",
    "StepRewardTransform",
    "OpenEnvRewardCalculator",
]