Spaces:
Runtime error
Runtime error
File size: 1,288 Bytes
a257639 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
from typing import Union
import gym
from stable_baselines3.common.callbacks import BaseCallback, StopTrainingOnNoModelImprovement
from stable_baselines3.common.vec_env import VecEnv, DummyVecEnv
class UpdateValEnv(BaseCallback):
def __init__(self, val_env: Union[gym.Env, VecEnv], verbose: int = 0):
super().__init__(verbose=verbose)
if not isinstance(val_env, VecEnv):
val_env = DummyVecEnv([lambda: val_env])
self.val_env = val_env
def _on_step(self) -> bool:
assert self.parent is not None, "``UpdateValEnv`` callback must be used with an ``EvalCallback``"
self.val_env.env_method('reset_instance_id')
return True
class UpdateValEnvAndStopTrainingOnNoModelImprovement(StopTrainingOnNoModelImprovement):
def __init__(self, val_env: Union[gym.Env, VecEnv],
max_no_improvement_evals: int, min_evals: int = 0, verbose: int = 0):
super().__init__(max_no_improvement_evals=max_no_improvement_evals, min_evals=min_evals, verbose=verbose)
if not isinstance(val_env, VecEnv):
val_env = DummyVecEnv([lambda: val_env])
self.val_env = val_env
def _on_step(self) -> bool:
self.val_env.env_method('reset_instance_id')
return super()._on_step()
|