Dominik Jain 3691ed2abc Support obs_rms persistence for MuJoCo by adding a general mechanism
for attaching persistence to Environments instances
2023-10-18 20:44:17 +02:00

33 lines
879 B
Python

import os
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING:
from tianshou.data import Collector
from tianshou.highlevel.env import Environments
from tianshou.highlevel.logger import Logger
from tianshou.policy import BasePolicy
from tianshou.trainer import BaseTrainer
@dataclass
class World:
envs: "Environments"
policy: "BasePolicy"
train_collector: "Collector"
test_collector: "Collector"
logger: "Logger"
restore_directory: str
trainer: Optional["BaseTrainer"] = None
@property
def persist_directory(self):
return self.logger.log_path
def persist_path(self, filename: str) -> str:
return os.path.join(self.persist_directory, filename)
def restore_path(self, filename: str) -> str:
return os.path.join(self.restore_directory, filename)