def train(env):
""" unit test helper function """
n_episodes = 10
n_steps_per_episode = 5
# setup logging
writer = LogStatsWriterTensorboard(log_dir='test_log',
tensorboard_render_figure=True)
register_log_stats_writer(writer)
# attach a console writer as well for immediate console feedback
register_log_stats_writer(LogStatsWriterConsole())
env = LogStatsWrapper.wrap(env, logging_prefix="train")
with SimpleStatsLoggingSetup(env):
for episode in range(n_episodes):
_ = env.reset()
for step in range(n_steps_per_episode):
# take random action
action = env.action_space.sample()
# take step in env and trigger log stats writing
_, _, _, _ = env.step(action)
# test accessing stats
env.get_stats(LogStatsLevel.EPOCH)
env.get_stats_value(BaseEnvEvents.reward, LogStatsLevel.EPOCH, name="mean")
def _launch_workers(self, env: ConfigType, wrappers: CollectionOfConfigType, agent: ConfigType) \
-> Iterable[Process]:
"""Configure the workers according to the rollout config and launch them."""
# Split total episode count across workers
episodes_per_process = [0] * self.n_processes
for i in range(self.n_episodes):
episodes_per_process[i % self.n_processes] += 1
# Configure and launch the processes
self.reporting_queue = Queue()
workers = []
for n_process_episodes in episodes_per_process:
if n_process_episodes == 0:
break
p = Process(
target=ParallelRolloutWorker.run,
args=(env, wrappers, agent,
n_process_episodes, self.max_episode_steps,
self.record_trajectory, self.input_dir, self.reporting_queue,
self.maze_seeding.generate_env_instance_seed(),
self.maze_seeding.generate_agent_instance_seed()),
daemon=True
)
p.start()
workers.append(p)
# Perform writer registration -- after the forks so that it is not carried over to child processes
if self.record_event_logs:
LogEventsWriterRegistry.register_writer(LogEventsWriterTSV(log_dir="./event_logs"))
register_log_stats_writer(LogStatsWriterConsole())
self.epoch_stats_aggregator = LogStatsAggregator(LogStatsLevel.EPOCH)
self.epoch_stats_aggregator.register_consumer(get_stats_logger("rollout_stats"))
return workers
def run_with(self, env: ConfigType, wrappers: CollectionOfConfigType,
agent: ConfigType):
"""Run the rollout sequentially in the main process."""
env, agent = self.init_env_and_agent(
env, wrappers, self.max_episode_steps, agent, self.input_dir,
self.maze_seeding.generate_env_instance_seed(),
self.maze_seeding.generate_agent_instance_seed())
# Set up the wrappers
# Hydra handles working directory
register_log_stats_writer(LogStatsWriterConsole())
if not isinstance(env, LogStatsWrapper):
env = LogStatsWrapper.wrap(env, logging_prefix="rollout_data")
if self.record_event_logs:
LogEventsWriterRegistry.register_writer(
LogEventsWriterTSV(log_dir="./event_logs"))
if self.record_trajectory:
TrajectoryWriterRegistry.register_writer(
TrajectoryWriterFile(log_dir="./trajectory_data"))
if not isinstance(env, TrajectoryRecordingWrapper):
env = TrajectoryRecordingWrapper.wrap(env)
self.progress_bar = tqdm(desc="Episodes done",
unit=" episodes",
total=self.n_episodes)
RolloutRunner.run_interaction_loop(
env,
agent,
self.n_episodes,
render=self.render,
episode_end_callback=lambda: self.update_progress())
self.progress_bar.close()
env.write_epoch_stats()
def __enter__(self) -> None:
"""Register a episode statistics aggregator with the provided log stats environment."""
# step logging setup: write to console
register_log_stats_writer(LogStatsWriterConsole())
if self.log_dir is not None:
register_log_stats_writer(
LogStatsWriterTensorboard(log_dir=self.log_dir,
tensorboard_render_figure=True))
def __enter__(self) -> T:
"""Register data writers."""
# Register stats, events, and trajectory data writers
print("*******", self.log_dir)
register_log_stats_writer(LogStatsWriterConsole())
register_log_stats_writer(
LogStatsWriterTensorboard(log_dir=self.log_dir + "/stats",
tensorboard_render_figure=True))
LogEventsWriterRegistry.register_writer(
LogEventsWriterTSV(log_dir=self.log_dir + "/event_logs"))
TrajectoryWriterRegistry.register_writer(
TrajectoryWriterFile(log_dir=self.log_dir + "/trajectory_data"))
return self.env
def setup_logging(job_config: Union[DictConfig, str, None],
log_dir: str = ".") -> None:
"""Setup tensorboard logging, derive the logging directory from the script name.
:param job_config: Configuration written as text to tensorboard (experiment config).
:param log_dir: log_dir for TensorBoard.
"""
# hydra handles the working directory
writer = LogStatsWriterTensorboard(log_dir=log_dir,
tensorboard_render_figure=True)
register_log_stats_writer(writer)
# attach a console writer as well for immediate console feedback
register_log_stats_writer(LogStatsWriterConsole())
summary_writer = writer.summary_writer
summary_writer.add_text("cmd", " ".join(sys.argv))
if job_config is not None:
# log run settings
if isinstance(job_config, DictConfig):
if job_config.__dict__["_metadata"].flags.get(
"allow_objects", False):
# Hydra was instructed to allow objects. This was done by the Python training API, hence we might have
# Python objects in our config, which makes logging the config to file not possible.
# todo Making this work for the Python training API would require reversing the Hydra instantiaton, i.e.
# generate a Hydra configuration from Python objects (at least partially).
BColors.print_colored(
"Logging run configurations with injected non-primitives is not supported yet. For now please don't"
" inject non-primitives if you wish to log the configuration of your run.",
BColors.WARNING)
return
else:
job_config = OmegaConf.to_yaml(job_config)
# prepare config text for tensorboard
job_config = job_config.replace("\n", "</br>")
job_config = job_config.replace(" ", " ")
summary_writer.add_text("job_config", job_config)
# Load the figures from the given files and add them to tensorboard.
for net_image_path in glob.glob('*.figure.pkl'):
network_description = net_image_path.split('/')[-1].replace(
'.figure.pkl', '')
fig = pickle.load(open(net_image_path, 'rb'))
summary_writer.add_figure(f'{network_description}', fig, close=True)
os.remove(net_image_path)
def init_logging(self, trainer_config: Dict[str, Any]) -> None:
"""Initialize logging.
This needs to be done here as the on_train_result is not called in the main process, but in a worker process.
Relies on the following:
- There should be only one Callbacks object per worker process
- The on train result should always be called in the same worker (if this is not the case,
this system should still handle it, but it might mess things up)
"""
assert self.epoch_stats is None, "Init logging should be called only once"
# Local epoch stats -- stats from all envs will be collected here together
self.epoch_stats = LogStatsAggregator(LogStatsLevel.EPOCH)
self.epoch_stats.register_consumer(get_stats_logger("train"))
# Initialize Tensorboard and console writers
writer = LogStatsWriterTensorboard(log_dir='.',
tensorboard_render_figure=True)
register_log_stats_writer(writer)
register_log_stats_writer(LogStatsWriterConsole())
summary_writer = writer.summary_writer
# Add config to tensorboard
yaml_config = pprint.pformat(trainer_config)
# prepare config text for tensorboard
yaml_config = yaml_config.replace("\n", "</br>")
yaml_config = yaml_config.replace(" ", " ")
summary_writer.add_text("job_config", yaml_config)
# Load the figures from the given files and add them to tensorboard.
network_files = filter(lambda x: x.endswith('.figure.pkl'),
os.listdir('.'))
for network_path in network_files:
network_name = network_path.split('/')[-1].replace(
'.figure.pkl', '')
fig = pickle.load(open(network_path, 'rb'))
summary_writer.add_figure(f'{network_name}', fig, close=True)
os.remove(network_path)
def test_observation_statistics_logging():
""" observation normalization logging test """
# normalization config
normalization_config = {
"default_strategy":
"maze.normalization_strategies.MeanZeroStdOneObservationNormalizationStrategy",
"default_strategy_config": {
"clip_range": (None, None),
"axis": 0
},
"default_statistics": None,
"statistics_dump": "statistics.pkl",
"exclude": None,
"manual_config": {
"observation": {
"strategy":
"maze.normalization_strategies.MeanZeroStdOneObservationNormalizationStrategy",
"strategy_config": {
"clip_range": (0, 1)
},
"statistics": {
"mean": [0, 0, 0, 0],
"std": [1, 1, 1, 1]
}
}
}
}
writer = LogStatsWriterTensorboard(log_dir='test_log',
tensorboard_render_figure=True)
register_log_stats_writer(writer)
# attach a console writer as well for immediate console feedback
register_log_stats_writer(LogStatsWriterConsole())
# init environment
env = GymMazeEnv("CartPole-v0")
# wrap env with observation normalization
env = ObservationNormalizationWrapper(
env,
default_strategy=normalization_config["default_strategy"],
default_strategy_config=normalization_config[
"default_strategy_config"],
default_statistics=normalization_config["default_statistics"],
statistics_dump=normalization_config["statistics_dump"],
sampling_policy=RandomPolicy(env.action_spaces_dict),
exclude=normalization_config["exclude"],
manual_config=normalization_config["manual_config"])
env = LogStatsWrapper.wrap(env, logging_prefix="train")
n_episodes = 10
n_steps_per_episode = 100
for episode in range(n_episodes):
_ = env.reset()
for step in range(n_steps_per_episode):
# take random action
action = env.action_space.sample()
# take step in env and trigger log stats writing
_, _, done, _ = env.step(action)
if done:
break
increment_log_step()