class Subscriber(ABC):
"""Event aggregation object.
"""
def __init__(self):
self.events = EventCollection()
self.reset()
@abstractmethod
def get_interfaces(self) -> List[Type[ABC]]:
"""
Specification of the event interfaces this subscriber wants to receive events from.
Every subscriber must implement this configuration method.
:return: A list of interface classes
"""
def notify_event(self, event: EventRecord):
"""Notify the subscriber of a new event occurrence.
:param event: the event
:return: None
"""
self.events.append(event)
def reset(self):
"""Reset event aggregation.
:return: None
"""
self.events = EventCollection()
def query_events(self, event_spec: Union[Callable, Iterable[Callable]]) -> Iterable:
"""Return all events collected at the current env step matching one or more given event types. The event
types are specified by the interface member function object itself.
Event calls are recorded as EventRecord, an object providing access to the passed arguments of the event
method.
:param event_spec: Specifies the event type by the interface member function. Can either be a single event
type specification or a list of specifications.
:return: An iterable to the event objects.
"""
# normalize event specification, convert to a list with a single item if necessary
if not isinstance(event_spec, Iterable):
event_spec = [event_spec]
for ev in event_spec:
assert self._get_class_that_defined_method(ev) in self.get_interfaces(), \
f"Event {ev} queried, but class not subscribed. Check your get_interfaces() implementation " \
f"{self.get_interfaces()}."
return self.events.query_events(event_spec)
@staticmethod
def _get_class_that_defined_method(method):
"""Source: https://stackoverflow.com/a/25959545"""
if inspect.ismethod(method):
for cls in inspect.getmro(method.__self__.__class__):
if cls.__dict__.get(method.__name__) is method:
return cls
method = method.__func__ # fallback to __qualname__ parsing
if inspect.isfunction(method):
cls = getattr(inspect.getmodule(method),
method.__qualname__.split('.<locals>', 1)[0].rsplit('.', 1)[0])
if isinstance(cls, type):
return cls
return getattr(method, '__objclass__', None) # handle special descriptor objects
class LogStatsWrapper(Wrapper[MazeEnv], LogStatsEnv):
"""A statistics logging wrapper for :class:`~maze.core.env.base_env.BaseEnv`.
:param env: The environment to wrap.
"""
def __init__(self, env: MazeEnv, logging_prefix: Optional[str] = None):
"""Avoid calling this constructor directly, use :method:`wrap` instead."""
# BaseEnv is a subset of gym.Env
super().__init__(env)
# initialize step aggregator
self.epoch_stats = LogStatsAggregator(LogStatsLevel.EPOCH)
self.episode_stats = LogStatsAggregator(LogStatsLevel.EPISODE, self.epoch_stats)
self.step_stats = LogStatsAggregator(LogStatsLevel.STEP, self.episode_stats)
self.stats_map = {
LogStatsLevel.EPOCH: self.epoch_stats,
LogStatsLevel.EPISODE: self.episode_stats,
LogStatsLevel.STEP: self.step_stats
}
if logging_prefix is not None:
self.epoch_stats.register_consumer(get_stats_logger(logging_prefix))
self.last_env_time: Optional[int] = None
self.reward_events = EventCollection()
self.episode_event_log: Optional[EpisodeEventLog] = None
self.step_stats_renderer = EventStatsRenderer()
# register a post-step callback, so stats are recorded even in case that a wrapper
# in the middle of the stack steps the environment (as done e.g. during step-skipping)
if hasattr(env, "context") and isinstance(env.context, EnvironmentContext):
env.context.register_post_step(self._record_stats_if_ready)
T = TypeVar("T")
@classmethod
def wrap(cls, env: T, logging_prefix: Optional[str] = None) -> Union[T, LogStatsEnv]:
"""Creation method providing appropriate type hints. Preferred method to construct the wrapper
compared to calling the class constructor directly.
:param env: The environment to be wrapped.
:param logging_prefix: The episode statistics is connected to the logging system with this tagging
prefix. If None, no logging happens.
:return A newly created wrapper instance.
"""
return cls(env, logging_prefix)
@override(BaseEnv)
def step(self, action: Any) -> Tuple[Any, Any, bool, Dict[Any, Any]]:
"""Collect the rewards for the logging statistics
"""
# get identifier of current substep
substep_id, _ = self.env.actor_id() if isinstance(self.env, StructuredEnv) else (None, None)
# take core env step
obs, rew, done, info = self.env.step(action)
# record the reward
self.reward_events.append(EventRecord(BaseEnvEvents, BaseEnvEvents.reward, dict(value=rew)))
self._record_stats_if_ready()
return obs, rew, done, info
def _record_stats_if_ready(self) -> None:
"""Checks if stats are ready to record based on env time (for structured envs, we wait till the end
of the whole structured step) and if so, does the recording.
"""
if self.last_env_time is None:
self.last_env_time = self.env.initial_env_time
# Recording of event logs and stats happens:
# - for TimeEnvs: Only if the env time changed, so that we record once per time step
# - for other envs: Every step
if isinstance(self.env, TimeEnvMixin) and self.env.get_env_time() == self.last_env_time:
return
step_event_log = StepEventLog(env_time=self.last_env_time, events=self.reward_events)
self.reward_events = EventCollection()
if isinstance(self.env, EventEnvMixin):
step_event_log.extend(self.env.get_step_events())
# add all recorded events to the step aggregator
for event_record in step_event_log.events:
self.step_stats.add_event(event_record)
# trigger logging statistics calculation
self.step_stats.reduce()
# lazy init new episode event log if needed
if not self.episode_event_log:
episode_id = self.env.get_episode_id() if isinstance(self.env, RecordableEnvMixin) else str(uuid.uuid4())
self.episode_event_log = EpisodeEventLog(episode_id)
# log raw events and init new step log
self.episode_event_log.step_event_logs.append(step_event_log)
# update the time of last stats recording
self.last_env_time = self.env.get_env_time() if isinstance(self.env, TimeEnvMixin) else self.last_env_time + 1
@override(BaseEnv)
def reset(self) -> Any:
"""Reset the environment and trigger the episode statistics calculation of the previous run.
"""
# Generate the episode stats from the previous rollout if any
self._calculate_kpis()
self.episode_stats.reduce()
self._write_episode_event_log()
# Initialize recording for the new episode (so we can record events already during env reset)
self.last_env_time = None
self.reward_events = EventCollection()
return self.env.reset()
@override(BaseEnv)
def close(self):
"""Close the stats rendering figure if needed."""
self.step_stats_renderer.close()
@override(LogStatsEnv)
def get_stats(self, level: LogStatsLevel) -> LogStatsAggregator:
"""Implementation of the LogStatsEnv interface, return the statistics aggregator."""
aggregator = self.stats_map[level]
return aggregator
@override(LogStatsEnv)
def write_epoch_stats(self):
"""Implementation of the LogStatsEnv interface, call reduce on the episode aggregator.
"""
if self.episode_event_log:
self._calculate_kpis()
self.episode_stats.reduce()
self.epoch_stats.reduce()
self._write_episode_event_log()
self.episode_event_log = None
@override(LogStatsEnv)
def get_stats_value(self,
event: Callable,
level: LogStatsLevel,
name: Optional[str] = None) -> LogStatsValue:
"""Implementation of the LogStatsEnv interface, obtain the value from the cached aggregator statistics.
"""
return self.epoch_stats.last_stats[(event, name, None)]
@override(LogStatsEnv)
def clear_epoch_stats(self) -> None:
"""Implementation of the LogStatsEnv interface, clear out episode statistics collected so far in this epoch."""
self.epoch_stats.clear_inputs()
def render_stats(self,
event_name: str = "BaseEnvEvents.reward",
metric_name: str = "value",
aggregation_func: Optional[Union[str, Callable]] = None,
group_by: str = None,
post_processing_func: Optional[Union[str, Callable]] = 'cumsum'):
"""Render statistics from the currently running episode.
Rendering is based on event logs. You can select arbitrary events from those dispatched by the currently
running environment.
:param event_name: Name of the even the even log corresponds to
:param metric_name: Metric to use (one of the event attributes, e.g. "n_items" -- depends on the event type)
:param aggregation_func: Optionally, specifies how to aggregate the metric on step level, i.e. when there
are multiple same events dispatched during the same step.
:param group_by: Optionally, another of event attributes to group by on the step level (e.g. "product_id")
:param post_processing_func: Optionally, a function to post-process the data ("cumsum" is often used)"""
self.step_stats_renderer.render_current_episode_stats(
self.episode_event_log, event_name, metric_name,
aggregation_func, group_by, post_processing_func)
def _calculate_kpis(self):
"""Calculate KPIs and append them to both aggregated and logged events."""
if not isinstance(self.env, EventEnvMixin) or not self.episode_event_log:
return
kpi_calculator = self.env.get_kpi_calculator()
if kpi_calculator is None:
return
last_maze_state = self.env.get_maze_state() if isinstance(self.env, RecordableEnvMixin) else None
kpis_dict = kpi_calculator.calculate_kpis(self.episode_event_log, last_maze_state)
kpi_events = []
for name, value in kpis_dict.items():
kpi_events.append(EventRecord(BaseEnvEvents, BaseEnvEvents.kpi, dict(name=name, value=value)))
for event_record in kpi_events:
self.episode_stats.add_event(event_record) # Add the events to episode aggregator
self.episode_event_log.step_event_logs[-1].events.append(event_record) # Log the events
def _write_episode_event_log(self):
"""Send the episode event log to writers."""
if self.episode_event_log:
LogEventsWriterRegistry.record_event_logs(self.episode_event_log)
self.episode_event_log = None
@override(Wrapper)
def get_observation_and_action_dicts(self, maze_state: Optional[MazeStateType],
maze_action: Optional[MazeActionType],
first_step_in_episode: bool) \
-> Tuple[Optional[Dict[Union[int, str], Any]], Optional[Dict[Union[int, str], Any]]]:
"""Keep both actions and observation the same."""
return self.env.get_observation_and_action_dicts(maze_state, maze_action, first_step_in_episode)
@override(SimulatedEnvMixin)
def clone_from(self, env: 'LogStatsWrapper') -> None:
"""implementation of :class:`~maze.core.env.simulated_env_mixin.SimulatedEnvMixin`."""
raise RuntimeError("Cloning the 'LogStatsWrapper' is not supported.")
def get_last_step_events(self, query: Union[Callable, Iterable[Callable]] = None):
"""Convenience accessor to all events recorded during the last step.
:param query: Specify which events to return (one or more interface methods)
:return: Recorded events from the last step (all if no query is present)
"""
if not self.episode_event_log or len(self.episode_event_log.step_event_logs) == 0:
return []
last_step_log = self.episode_event_log.step_event_logs[-1]
if query:
return list(last_step_log.events.query_events(query))
else:
return list(last_step_log.events.events)