app/training/rl_dataset.py

"""RL episode logging dataset."""

from __future__ import annotations


def make_rl_record(observation: dict, action: dict, reward: float, done: bool) -> dict:
    return {
        "observation": observation,
        "action": action,
        "reward": reward,
        "done": done,
    }