Skip to content

Allow injecting reward and controller classes into AEES #22

Description

@davidkfoss

Problem

AEES.__init__ hardcodes the reward and controller construction. Even though BaseReward and BaseController are public protocols, every user gets:

  • NormalizedLossImprovementReward
  • DiscountedUCBController / BucketedContextualController / RandomController

There is no injection point, so the protocol abstraction is half-finished — the types are public but the instantiation is closed.

Proposal

Accept optional pre-built controller / reward instances (or factories) on AEES.__init__, falling back to today's defaults when not provided. Something like:

AEES(
    optimizer,
    reward: BaseReward | None = None,
    lr_controller: BaseController | None = None,
    noise_controller: BaseController | list[BaseController] | None = None,
    ...
)

Acceptance criteria

  • Users can pass custom BaseReward / BaseController implementations.
  • Defaults remain unchanged when nothing is passed (no behavior regression).
  • Validation still rejects mismatched controller / context-mode combinations.
  • Docs/example updated to show how to plug in a custom controller.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions