Class CsvWriter

qlib/rl/utils/log.py:470–508

Dump all episode metrics to a ``result.csv``. This is not the correct implementation. It's only used for first iteration.

class CsvWriter(LogWriter):
    """Dump all episode metrics to a ``result.csv``.

    This is not the correct implementation. It's only used for first iteration.
    """

    SUPPORTED_TYPES = (float, str, pd.Timestamp)

    all_records: List[Dict[str, Any]]

    # FIXME: save & reload

    def __init__(self, output_dir: Path, loglevel: int | LogLevel = LogLevel.PERIODIC) -> None:
        super().__init__(loglevel)
        self.output_dir = output_dir
        self.output_dir.mkdir(exist_ok=True)

    def clear(self) -> None:
        super().clear()
        self.all_records = []

    def log_episode(self, length: int, rewards: List[float], contents: List[Dict[str, Any]]) -> None:
        # FIXME Same as ConsoleLogger, needs a refactor to eliminate code-dup
        episode_wise_contents: Dict[str, list] = defaultdict(list)

        for step_contents in contents:
            for name, value in step_contents.items():
                if isinstance(value, self.SUPPORTED_TYPES):
                    episode_wise_contents[name].append(value)

        logs: Dict[str, float] = {}
        for name, values in episode_wise_contents.items():
            logs[name] = self.aggregation(values, name)  # type: ignore

        self.all_records.append(logs)

    def on_env_all_done(self) -> None:
        # FIXME: this is temporary
        pd.DataFrame.from_records(self.all_records).to_csv(self.output_dir / "result.csv", index=False)


# The following are not implemented yet.