"""
Data types for neuro-evolution training
"""
import chex


@chex.dataclass
class TrainingData:
    """
    Dataclass tracking training rewards and recorded data

    Simulations used for training should return
    a ``TrainingData`` class as its recorded data.

    Examples
    --------

    .. code-block:: python

       def step(self, i, k, params, state):
           ...
           # Rewards are returned as part of the TrainingData
           return (
               state,
               TrainingData(rewards=rewards, records=records)
           )

    Parameters
    ----------
    rewards
        PyTree of rewards generated by simulation agents.
    records
        Any additional data recorded over the course
        of the simulation.
    """

    rewards: chex.Array
    records: chex.ArrayTree
