Source code for unilab.envs.mdp.recorders

"""Reusable recorder terms for the Manager-Based tensor runtime."""

from __future__ import annotations

from typing import TYPE_CHECKING

import torch

from unilab.managers.recorder_manager import RecorderTerm, RecorderTermCfg

if TYPE_CHECKING:
    from unilab.managers.types import ManagerBasedRlEnv




class LifecycleCounterRecorder(RecorderTerm):
    """Count recorder lifecycle calls without performing I/O."""



    def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedRlEnv):
        super().__init__(cfg, env)
        self.pre_reset_count = torch.zeros((), dtype=torch.int64)
        self.post_reset_count = torch.zeros((), dtype=torch.int64)
        self.post_step_count = 0




    def record_pre_reset(self, env_ids: torch.Tensor) -> None:
        self.pre_reset_count += env_ids.numel()




    def record_post_reset(self, env_ids: torch.Tensor) -> None:
        self.post_reset_count += env_ids.numel()




    def record_post_step(self) -> None:
        self.post_step_count += 1




__all__ = ["LifecycleCounterRecorder"]