Source code for metasim.cfg.checkers.checker_operators

from __future__ import annotations

from dataclasses import MISSING

import torch

from metasim.cfg.objects import BaseObjMetaCfg
from metasim.utils.configclass import configclass

from .base_checker import BaseChecker

try:
    from metasim.sim.base import BaseSimHandler
except:
    pass




@configclass
class AndOp(BaseChecker):
    checkers: list[BaseChecker] = MISSING



    def reset(self, handler: BaseSimHandler, env_ids: list[int] | None = None):
        for checker in self.checkers:
            checker.reset(handler, env_ids=env_ids)




    def check(self, handler: BaseSimHandler) -> torch.BoolTensor:
        success = torch.ones(handler.num_envs, dtype=torch.bool)
        for checker in self.checkers:
            success = success & checker.check(handler)
        return success




    def get_debug_viewers(self) -> list[BaseObjMetaCfg]:
        viewers = []
        for checker in self.checkers:
            viewers += checker.get_debug_viewers()
        return viewers






@configclass
class OrOp(BaseChecker):
    checkers: list[BaseChecker] = MISSING



    def reset(self, handler: BaseSimHandler, env_ids: list[int] | None = None):
        for checker in self.checkers:
            checker.reset(handler, env_ids=env_ids)




    def check(self, handler: BaseSimHandler) -> torch.BoolTensor:
        success = torch.zeros(handler.num_envs, dtype=torch.bool)
        for checker in self.checkers:
            success = success | checker.check(handler)
        return success




    def get_debug_viewers(self) -> list[BaseObjMetaCfg]:
        viewers = []
        for checker in self.checkers:
            viewers += checker.get_debug_viewers()
        return viewers






@configclass
class NotOp(BaseChecker):
    checker: BaseChecker = MISSING



    def reset(self, handler: BaseSimHandler, env_ids: list[int] | None = None):
        self.checker.reset(handler, env_ids=env_ids)




    def check(self, handler: BaseSimHandler) -> torch.BoolTensor:
        return ~self.checker.check(handler)




    def get_debug_viewers(self) -> list[BaseObjMetaCfg]:
        return self.checker.get_debug_viewers()