Shortcuts

Source code for ding.worker.collector.sample_serial_collector

from typing import Optional, Any, List
from collections import namedtuple
from easydict import EasyDict
import copy
import numpy as np
import torch

from ding.envs import BaseEnvManager
from ding.utils import build_logger, EasyTimer, SERIAL_COLLECTOR_REGISTRY, one_time_warning, get_rank, get_world_size, \
    allreduce_data
from ding.torch_utils import to_tensor, to_ndarray
from .base_serial_collector import ISerialCollector, CachePool, TrajBuffer, INF, to_tensor_transitions


[docs]@SERIAL_COLLECTOR_REGISTRY.register('sample') class SampleSerialCollector(ISerialCollector): """ Overview: Sample collector(n_sample), a sample is one training sample for updating model, it is usually like <s, a, s', r, d>(one transition) while is a trajectory with many transitions, which is often used in RNN-model. Interfaces: __init__, reset, reset_env, reset_policy, collect, close Property: envstep """ config = dict(deepcopy_obs=False, transform_obs=False, collect_print_freq=100)
[docs] def __init__( self, cfg: EasyDict, env: BaseEnvManager = None, policy: namedtuple = None, tb_logger: 'SummaryWriter' = None, # noqa exp_name: Optional[str] = 'default_experiment', instance_name: Optional[str] = 'collector' ) -> None: """ Overview: Initialization method. Arguments: - cfg (:obj:`EasyDict`): Config dict - env (:obj:`BaseEnvManager`): the subclass of vectorized env_manager(BaseEnvManager) - policy (:obj:`namedtuple`): the api namedtuple of collect_mode policy - tb_logger (:obj:`SummaryWriter`): tensorboard handle """ self._exp_name = exp_name self._instance_name = instance_name self._collect_print_freq = cfg.collect_print_freq self._deepcopy_obs = cfg.deepcopy_obs # whether to deepcopy each data self._transform_obs = cfg.transform_obs self._cfg = cfg self._timer = EasyTimer() self._end_flag = False self._rank = get_rank() self._world_size = get_world_size() if self._rank == 0: if tb_logger is not None: self._logger, _ = build_logger( path='./{}/log/{}'.format(self._exp_name, self._instance_name), name=self._instance_name, need_tb=False ) self._tb_logger = tb_logger else: self._logger, self._tb_logger = build_logger( path='./{}/log/{}'.format(self._exp_name, self._instance_name), name=self._instance_name ) else: self._logger, _ = build_logger( path='./{}/log/{}'.format(self._exp_name, self._instance_name), name=self._instance_name, need_tb=False ) self._tb_logger = None self.reset(policy, env)
[docs] def reset_env(self, _env: Optional[BaseEnvManager] = None) -> None: """ Overview: Reset the environment. If _env is None, reset the old environment. If _env is not None, replace the old environment in the collector with the new passed \ in environment and launch. Arguments: - env (:obj:`Optional[BaseEnvManager]`): instance of the subclass of vectorized \ env_manager(BaseEnvManager) """ if _env is not None: self._env = _env self._env.launch() self._env_num = self._env.env_num else: self._env.reset()
[docs] def reset_policy(self, _policy: Optional[namedtuple] = None) -> None: """ Overview: Reset the policy. If _policy is None, reset the old policy. If _policy is not None, replace the old policy in the collector with the new passed in policy. Arguments: - policy (:obj:`Optional[namedtuple]`): the api namedtuple of collect_mode policy """ assert hasattr(self, '_env'), "please set env first" if _policy is not None: self._policy = _policy self._policy_cfg = self._policy.get_attribute('cfg') self._default_n_sample = _policy.get_attribute('n_sample') self._traj_len_inf = self._policy_cfg.traj_len_inf self._unroll_len = _policy.get_attribute('unroll_len') self._on_policy = _policy.get_attribute('on_policy') if self._default_n_sample is not None and not self._traj_len_inf: self._traj_len = max( self._unroll_len, self._default_n_sample // self._env_num + int(self._default_n_sample % self._env_num != 0) ) self._logger.debug( 'Set default n_sample mode(n_sample({}), env_num({}), traj_len({}))'.format( self._default_n_sample, self._env_num, self._traj_len ) ) else: self._traj_len = INF self._policy.reset()
[docs] def reset(self, _policy: Optional[namedtuple] = None, _env: Optional[BaseEnvManager] = None) -> None: """ Overview: Reset the environment and policy. If _env is None, reset the old environment. If _env is not None, replace the old environment in the collector with the new passed \ in environment and launch. If _policy is None, reset the old policy. If _policy is not None, replace the old policy in the collector with the new passed in policy. Arguments: - policy (:obj:`Optional[namedtuple]`): the api namedtuple of collect_mode policy - env (:obj:`Optional[BaseEnvManager]`): instance of the subclass of vectorized \ env_manager(BaseEnvManager) """ if _env is not None: self.reset_env(_env) if _policy is not None: self.reset_policy(_policy) if self._policy_cfg.type == 'dreamer_command': self._states = None self._resets = np.array([False for i in range(self._env_num)]) self._obs_pool = CachePool('obs', self._env_num, deepcopy=self._deepcopy_obs) self._policy_output_pool = CachePool('policy_output', self._env_num) # _traj_buffer is {env_id: TrajBuffer}, is used to store traj_len pieces of transitions maxlen = self._traj_len if self._traj_len != INF else None self._traj_buffer = { env_id: TrajBuffer(maxlen=maxlen, deepcopy=self._deepcopy_obs) for env_id in range(self._env_num) } self._env_info = {env_id: {'time': 0., 'step': 0, 'train_sample': 0} for env_id in range(self._env_num)} self._episode_info = [] self._total_envstep_count = 0 self._total_episode_count = 0 self._total_train_sample_count = 0 self._total_duration = 0 self._last_train_iter = 0 self._end_flag = False
[docs] def _reset_stat(self, env_id: int) -> None: """ Overview: Reset the collector's state. Including reset the traj_buffer, obs_pool, policy_output_pool\ and env_info. Reset these states according to env_id. You can refer to base_serial_collector\ to get more messages. Arguments: - env_id (:obj:`int`): the id where we need to reset the collector's state """ self._traj_buffer[env_id].clear() self._obs_pool.reset(env_id) self._policy_output_pool.reset(env_id) self._env_info[env_id] = {'time': 0., 'step': 0, 'train_sample': 0}
@property def envstep(self) -> int: """ Overview: Print the total envstep count. Return: - envstep (:obj:`int`): The total envstep count. """ return self._total_envstep_count @envstep.setter def envstep(self, value: int) -> None: """ Overview: Set the total envstep count. Arguments: - value (:obj:`int`): The total envstep count. """ self._total_envstep_count = value
[docs] def close(self) -> None: """ Overview: Close the collector. If end_flag is False, close the environment, flush the tb_logger\ and close the tb_logger. """ if self._end_flag: return self._end_flag = True self._env.close() if self._tb_logger: self._tb_logger.flush() self._tb_logger.close()
[docs] def __del__(self) -> None: """ Overview: Execute the close command and close the collector. __del__ is automatically called to \ destroy the collector instance when the collector finishes its work """ self.close()
[docs] def collect( self, n_sample: Optional[int] = None, train_iter: int = 0, drop_extra: bool = True, random_collect: bool = False, record_random_collect: bool = True, policy_kwargs: Optional[dict] = None, level_seeds: Optional[List] = None, ) -> List[Any]: """ Overview: Collect `n_sample` data with policy_kwargs, which is already trained `train_iter` iterations. Arguments: - n_sample (:obj:`int`): The number of collecting data sample. - train_iter (:obj:`int`): The number of training iteration when calling collect method. - drop_extra (:obj:`bool`): Whether to drop extra return_data more than `n_sample`. - record_random_collect (:obj:`bool`) :Whether to output logs of random collect. - policy_kwargs (:obj:`dict`): The keyword args for policy forward. - level_seeds (:obj:`dict`): Used in PLR, represents the seed of the environment that \ generate the data Returns: - return_data (:obj:`List`): A list containing training samples. """ if n_sample is None: if self._default_n_sample is None: raise RuntimeError("Please specify collect n_sample") else: n_sample = self._default_n_sample if n_sample % self._env_num != 0: one_time_warning( "Please make sure env_num is divisible by n_sample: {}/{}, ".format(n_sample, self._env_num) + "which may cause convergence problems in a few algorithms" ) if policy_kwargs is None: policy_kwargs = {} collected_sample = 0 collected_step = 0 collected_episode = 0 return_data = [] while collected_sample < n_sample: with self._timer: # Get current env obs. obs = self._env.ready_obs # Policy forward. self._obs_pool.update(obs) if self._transform_obs: obs = to_tensor(obs, dtype=torch.float32) if self._policy_cfg.type == 'dreamer_command' and not random_collect: policy_output = self._policy.forward(obs, **policy_kwargs, reset=self._resets, state=self._states) #self._states = {env_id: output['state'] for env_id, output in policy_output.items()} self._states = [output['state'] for output in policy_output.values()] else: policy_output = self._policy.forward(obs, **policy_kwargs) self._policy_output_pool.update(policy_output) # Interact with env. actions = {env_id: output['action'] for env_id, output in policy_output.items()} actions = to_ndarray(actions) timesteps = self._env.step(actions) # TODO(nyz) this duration may be inaccurate in async env interaction_duration = self._timer.value / len(timesteps) # TODO(nyz) vectorize this for loop for env_id, timestep in timesteps.items(): with self._timer: if timestep.info.get('abnormal', False): # If there is an abnormal timestep, reset all the related variables(including this env). # suppose there is no reset param, just reset this env self._env.reset({env_id: None}) self._policy.reset([env_id]) self._reset_stat(env_id) self._logger.info('Env{} returns a abnormal step, its info is {}'.format(env_id, timestep.info)) continue if self._policy_cfg.type == 'dreamer_command' and not random_collect: self._resets[env_id] = timestep.done if self._policy_cfg.type == 'ngu_command': # for NGU policy transition = self._policy.process_transition( self._obs_pool[env_id], self._policy_output_pool[env_id], timestep, env_id ) else: transition = self._policy.process_transition( self._obs_pool[env_id], self._policy_output_pool[env_id], timestep ) if level_seeds is not None: transition['seed'] = level_seeds[env_id] # ``train_iter`` passed in from ``serial_entry``, indicates current collecting model's iteration. transition['collect_iter'] = train_iter self._traj_buffer[env_id].append(transition) self._env_info[env_id]['step'] += 1 collected_step += 1 # prepare data if timestep.done or len(self._traj_buffer[env_id]) == self._traj_len: # If policy is r2d2: # 1. For each collect_env, we want to collect data of length self._traj_len=INF # unless the episode enters the 'done' state. # 2. The length of a train (sequence) sample in r2d2 is <burnin + learn_unroll_length> # (please refer to r2d2.py) and in each collect phase, # we collect a total of <n_sample> (sequence) samples. # 3. When timestep is done and we only collected very few transitions in self._traj_buffer, # by going through self._policy.get_train_sample, it will be padded automatically to get the # sequence sample of length <burnin + learn_unroll_len> (please refer to r2d2.py). # Episode is done or traj_buffer(maxlen=traj_len) is full. # indicate whether to shallow copy next obs, i.e., overlap of s_t and s_t+1 transitions = to_tensor_transitions(self._traj_buffer[env_id], not self._deepcopy_obs) train_sample = self._policy.get_train_sample(transitions) return_data.extend(train_sample) self._env_info[env_id]['train_sample'] += len(train_sample) collected_sample += len(train_sample) self._traj_buffer[env_id].clear() self._env_info[env_id]['time'] += self._timer.value + interaction_duration # If env is done, record episode info and reset if timestep.done: collected_episode += 1 reward = timestep.info['eval_episode_return'] info = { 'reward': reward, 'time': self._env_info[env_id]['time'], 'step': self._env_info[env_id]['step'], 'train_sample': self._env_info[env_id]['train_sample'], } self._episode_info.append(info) # Env reset is done by env_manager automatically self._policy.reset([env_id]) self._reset_stat(env_id) collected_duration = sum([d['time'] for d in self._episode_info]) # reduce data when enables DDP if self._world_size > 1: collected_sample = allreduce_data(collected_sample, 'sum') collected_step = allreduce_data(collected_step, 'sum') collected_episode = allreduce_data(collected_episode, 'sum') collected_duration = allreduce_data(collected_duration, 'sum') self._total_envstep_count += collected_step self._total_episode_count += collected_episode self._total_duration += collected_duration self._total_train_sample_count += collected_sample # log if record_random_collect: # default is true, but when random collect, record_random_collect is False self._output_log(train_iter) else: self._episode_info.clear() # on-policy reset if self._on_policy: for env_id in range(self._env_num): self._reset_stat(env_id) if drop_extra: return return_data[:n_sample] else: return return_data
[docs] def _output_log(self, train_iter: int) -> None: """ Overview: Print the output log information. You can refer to the docs of `Best Practice` to understand \ the training generated logs and tensorboards. Arguments: - train_iter (:obj:`int`): the number of training iteration. """ if self._rank != 0: return if (train_iter - self._last_train_iter) >= self._collect_print_freq and len(self._episode_info) > 0: self._last_train_iter = train_iter episode_count = len(self._episode_info) envstep_count = sum([d['step'] for d in self._episode_info]) train_sample_count = sum([d['train_sample'] for d in self._episode_info]) duration = sum([d['time'] for d in self._episode_info]) episode_return = [d['reward'] for d in self._episode_info] info = { 'episode_count': episode_count, 'envstep_count': envstep_count, 'train_sample_count': train_sample_count, 'avg_envstep_per_episode': envstep_count / episode_count, 'avg_sample_per_episode': train_sample_count / episode_count, 'avg_envstep_per_sec': envstep_count / duration, 'avg_train_sample_per_sec': train_sample_count / duration, 'avg_episode_per_sec': episode_count / duration, 'reward_mean': np.mean(episode_return), 'reward_std': np.std(episode_return), 'reward_max': np.max(episode_return), 'reward_min': np.min(episode_return), 'total_envstep_count': self._total_envstep_count, 'total_train_sample_count': self._total_train_sample_count, 'total_episode_count': self._total_episode_count, # 'each_reward': episode_return, } self._episode_info.clear() self._logger.info("collect end:\n{}".format('\n'.join(['{}: {}'.format(k, v) for k, v in info.items()]))) for k, v in info.items(): if k in ['each_reward']: continue self._tb_logger.add_scalar('{}_iter/'.format(self._instance_name) + k, v, train_iter) if k in ['total_envstep_count']: continue self._tb_logger.add_scalar('{}_step/'.format(self._instance_name) + k, v, self._total_envstep_count)