Source code for palaestrai.environment.environment_state

from __future__ import annotations

import dataclasses
from typing import TYPE_CHECKING, Any, List, Optional

if TYPE_CHECKING:
    from palaestrai.types import SimTime
    from palaestrai.agent import SensorInformation, RewardInformation




@dataclasses.dataclass
class EnvironmentState:
    """Describes the current state of an :class:`~Environment`.

    This dataclass is used as return value of the :meth:`~Environment.update()`
    method. It contains current sensor readings, reward of the environment,
    indicates whether the environment has terminated or not, and finally gives
    time information.

    Terminated vs. truncated
    ------------------------

    In addition to the coarse ``done`` flag (which merely says "the episode
    is over, please reset"), :class:`EnvironmentState` carries the more
    specific gymnasium-style pair :attr:`terminated` / :attr:`truncated`.
    They describe *why* the episode ended and are the pieces of information
    a learning algorithm needs in order to bootstrap the value function
    correctly at episode boundaries:

    * ``terminated=True`` means the environment reached a genuine terminal
      state of the underlying Markov Decision Process (MDP). No further
      transitions are defined; the Bellman backup at this step must not
      bootstrap (``target = r``). Examples: a game was won/lost, a robot
      crashed, a power-flow solver diverged and the grid collapsed.
    * ``truncated=True`` means the episode was cut off by something
      *outside* the MDP (a time limit, a scenario horizon, a training-loop
      budget). The state has a successor that simply wasn't observed;
      the Bellman backup must still bootstrap (``target = r + γ · V(s')``).
    * Both ``None`` (default) preserves the legacy contract: consumers
      that only look at ``done`` will keep working, and consumers that
      know about the new fields should fall back to treating ``done`` as
      termination. This makes the two fields opt-in per environment.

    Invariants (enforced in :meth:`__post_init__`):

    * ``terminated`` and ``truncated`` are mutually exclusive: at most one
      of them may be ``True`` for the same transition.
    * If either flag is ``True``, ``done`` must also be ``True`` — a
      terminal or truncated step is by definition an episode-ending step.

    Attributes
    ----------
    sensor_information : List[SensorInformation]
        List of current sensor values after evaluating the environment.
    rewards : List[RewardInformation]
        Current rewards given from the environment.
    done : bool
        Whether the environment has finished this episode (``True``) or
        not (``False``). This is the coarse "please reset" bit and is
        equivalent to ``terminated or truncated`` in the gymnasium API
        for environments that populate both.
    terminated : Optional[bool] (default: ``None``)
        ``True`` iff the episode ended in a genuine MDP terminal state.
        ``None`` means the environment does not yet distinguish
        terminated from truncated; consumers should fall back to the
        legacy behavior of treating ``done`` as termination.
    truncated : Optional[bool] (default: ``None``)
        ``True`` iff the episode ended for reasons outside the MDP (time
        limit, scenario horizon, user abort). ``None`` has the same
        meaning as for :attr:`terminated`.
    world_state : Any (default: ``None``)
        Current state of the world (whatever the environment thinks it
        is). Kept for backwards compatibility as an extension slot for
        environment-specific side channels.
    simtime : SimTime (default: ``None``)
        Environment starting time.
    """

    sensor_information: List[SensorInformation]
    rewards: List[RewardInformation]
    done: bool
    world_state: Any = None
    simtime: Optional[SimTime] = None
    # NOTE: `terminated` and `truncated` are intentionally placed at the
    # tail of the dataclass. Positional callers of the historical
    # signature (``EnvironmentState(sensors, rewards, done)`` or
    # ``EnvironmentState(sensors, rewards, done, world_state, simtime)``)
    # continue to work bit-for-bit; only new callers that opt into the
    # gymnasium-style distinction need to pass the two extra fields, and
    # they will typically do so by keyword.
    terminated: Optional[bool] = None
    truncated: Optional[bool] = None

    def __post_init__(self) -> None:
        # At most one of terminated/truncated may be True.
        if self.terminated is True and self.truncated is True:
            raise ValueError(
                "EnvironmentState: `terminated` and `truncated` are "
                "mutually exclusive; at most one may be True for the "
                "same transition."
            )
        # If either is True, `done` must be True.
        if (
            self.terminated is True or self.truncated is True
        ) and not self.done:
            raise ValueError(
                "EnvironmentState: `terminated=True` or `truncated=True` "
                "implies `done=True`; got done=False."
            )