import numpy as np from torch.utils.tensorboard import SummaryWriter from tianshou.policy import BasePolicy from tianshou.env import VectorEnv, SubprocVectorEnv, AsyncVectorEnv from tianshou.data import Collector, Batch, ReplayBuffer if __name__ == '__main__': from env import MyTestEnv else: # pytest from test.base.env import MyTestEnv class MyPolicy(BasePolicy): def __init__(self, dict_state=False): super().__init__() self.dict_state = dict_state def forward(self, batch, state=None): if self.dict_state: return Batch(act=np.ones(len(batch.obs['index']))) return Batch(act=np.ones(len(batch.obs))) def learn(self): pass class Logger: def __init__(self, writer): self.cnt = 0 self.writer = writer def preprocess_fn(self, **kwargs): # modify info before adding into the buffer, and recorded into tfb # if info is not provided from env, it will be a ``Batch()``. if not kwargs.get('info', Batch()).is_empty(): n = len(kwargs['obs']) info = kwargs['info'] for i in range(n): info[i].update(rew=kwargs['rew'][i]) self.writer.add_scalar('key', np.mean( info['key']), global_step=self.cnt) self.cnt += 1 return Batch(info=info) # or: return {'info': info} else: return Batch() @staticmethod def single_preprocess_fn(**kwargs): # same as above, without tfb if not kwargs.get('info', Batch()).is_empty(): n = len(kwargs['obs']) info = kwargs['info'] for i in range(n): info[i].update(rew=kwargs['rew'][i]) return Batch(info=info) # or: return {'info': info} else: return Batch() def test_collector(): writer = SummaryWriter('log/collector') logger = Logger(writer) env_fns = [lambda x=i: MyTestEnv(size=x, sleep=0) for i in [2, 3, 4, 5]] venv = SubprocVectorEnv(env_fns) dum = VectorEnv(env_fns) policy = MyPolicy() env = env_fns[0]() c0 = Collector(policy, env, ReplayBuffer(size=100, ignore_obs_next=False), logger.preprocess_fn) c0.collect(n_step=3) assert np.allclose(c0.buffer.obs[:4], np.expand_dims( [0, 1, 0, 1], axis=-1)) assert np.allclose(c0.buffer[:4].obs_next, np.expand_dims( [1, 2, 1, 2], axis=-1)) c0.collect(n_episode=3) assert np.allclose(c0.buffer.obs[:10], np.expand_dims( [0, 1, 0, 1, 0, 1, 0, 1, 0, 1], axis=-1)) assert np.allclose(c0.buffer[:10].obs_next, np.expand_dims( [1, 2, 1, 2, 1, 2, 1, 2, 1, 2], axis=-1)) c0.collect(n_step=3, random=True) c1 = Collector(policy, venv, ReplayBuffer(size=100, ignore_obs_next=False), logger.preprocess_fn) c1.collect(n_step=6) assert np.allclose(c1.buffer.obs[:11], np.expand_dims( [0, 1, 0, 1, 2, 0, 1, 0, 1, 2, 3], axis=-1)) assert np.allclose(c1.buffer[:11].obs_next, np.expand_dims([ 1, 2, 1, 2, 3, 1, 2, 1, 2, 3, 4], axis=-1)) c1.collect(n_episode=2) assert np.allclose(c1.buffer.obs[11:21], np.expand_dims( [0, 1, 2, 3, 4, 0, 1, 0, 1, 2], axis=-1)) assert np.allclose(c1.buffer[11:21].obs_next, np.expand_dims([1, 2, 3, 4, 5, 1, 2, 1, 2, 3], axis=-1)) c1.collect(n_episode=3, random=True) c2 = Collector(policy, dum, ReplayBuffer(size=100, ignore_obs_next=False), logger.preprocess_fn) c2.collect(n_episode=[1, 2, 2, 2]) assert np.allclose(c2.buffer.obs_next[:26], np.expand_dims([ 1, 2, 1, 2, 3, 1, 2, 3, 4, 1, 2, 3, 4, 5, 1, 2, 3, 1, 2, 3, 4, 1, 2, 3, 4, 5], axis=-1)) c2.reset_env() c2.collect(n_episode=[2, 2, 2, 2]) assert np.allclose(c2.buffer.obs_next[26:54], np.expand_dims([ 1, 2, 1, 2, 3, 1, 2, 1, 2, 3, 4, 1, 2, 3, 4, 5, 1, 2, 3, 1, 2, 3, 4, 1, 2, 3, 4, 5], axis=-1)) c2.collect(n_episode=[1, 1, 1, 1], random=True) def test_collector_with_async(): env_lens = [2, 3, 4, 5] writer = SummaryWriter('log/async_collector') logger = Logger(writer) env_fns = [lambda x=i: MyTestEnv(size=x, sleep=0.1, random_sleep=True) for i in env_lens] venv = AsyncVectorEnv(env_fns) policy = MyPolicy() c1 = Collector(policy, venv, ReplayBuffer(size=1000, ignore_obs_next=False), logger.preprocess_fn) c1.collect(n_episode=10) # check if the data in the buffer is chronological # i.e. data in the buffer are full episodes, and each episode is # returned by the same environment env_id = c1.buffer.info['env_id'] size = len(c1.buffer) obs = c1.buffer.obs[:size] done = c1.buffer.done[:size] print(env_id[:size]) print(obs) obs_ground_truth = [] i = 0 while i < size: # i is the start of an episode if done[i]: # this episode has one transition assert env_lens[env_id[i]] == 1 i += 1 continue j = i while True: j += 1 # in one episode, the environment id is the same assert env_id[j] == env_id[i] if done[j]: break j = j + 1 # j is the start of the next episode assert j - i == env_lens[env_id[i]] obs_ground_truth += list(range(j - i)) i = j obs_ground_truth = np.expand_dims( np.array(obs_ground_truth), axis=-1) assert np.allclose(obs, obs_ground_truth) def test_collector_with_dict_state(): env = MyTestEnv(size=5, sleep=0, dict_state=True) policy = MyPolicy(dict_state=True) c0 = Collector(policy, env, ReplayBuffer(size=100), Logger.single_preprocess_fn) c0.collect(n_step=3) c0.collect(n_episode=2) env_fns = [lambda x=i: MyTestEnv(size=x, sleep=0, dict_state=True) for i in [2, 3, 4, 5]] envs = VectorEnv(env_fns) envs.seed(666) obs = envs.reset() assert not np.isclose(obs[0]['rand'], obs[1]['rand']) c1 = Collector(policy, envs, ReplayBuffer(size=100), Logger.single_preprocess_fn) c1.seed(0) c1.collect(n_step=10) c1.collect(n_episode=[2, 1, 1, 2]) batch = c1.sample(10) print(batch) c0.buffer.update(c1.buffer) assert np.allclose(c0.buffer[:len(c0.buffer)].obs.index, np.expand_dims([ 0., 1., 2., 3., 4., 0., 1., 2., 3., 4., 0., 1., 2., 3., 4., 0., 1., 0., 1., 2., 0., 1., 0., 1., 2., 3., 0., 1., 2., 3., 4., 0., 1., 0., 1., 2., 0., 1., 0., 1., 2., 3., 0., 1., 2., 3., 4.], axis=-1)) c2 = Collector(policy, envs, ReplayBuffer(size=100, stack_num=4), Logger.single_preprocess_fn) c2.collect(n_episode=[0, 0, 0, 10]) batch = c2.sample(10) print(batch['obs_next']['index']) def test_collector_with_ma(): def reward_metric(x): return x.sum() env = MyTestEnv(size=5, sleep=0, ma_rew=4) policy = MyPolicy() c0 = Collector(policy, env, ReplayBuffer(size=100), Logger.single_preprocess_fn, reward_metric=reward_metric) # n_step=3 will collect a full episode r = c0.collect(n_step=3)['rew'] assert np.asanyarray(r).size == 1 and r == 4. r = c0.collect(n_episode=2)['rew'] assert np.asanyarray(r).size == 1 and r == 4. env_fns = [lambda x=i: MyTestEnv(size=x, sleep=0, ma_rew=4) for i in [2, 3, 4, 5]] envs = VectorEnv(env_fns) c1 = Collector(policy, envs, ReplayBuffer(size=100), Logger.single_preprocess_fn, reward_metric=reward_metric) r = c1.collect(n_step=10)['rew'] assert np.asanyarray(r).size == 1 and r == 4. r = c1.collect(n_episode=[2, 1, 1, 2])['rew'] assert np.asanyarray(r).size == 1 and r == 4. batch = c1.sample(10) print(batch) c0.buffer.update(c1.buffer) obs = np.array(np.expand_dims([ 0., 1., 2., 3., 4., 0., 1., 2., 3., 4., 0., 1., 2., 3., 4., 0., 1., 0., 1., 2., 0., 1., 0., 1., 2., 3., 0., 1., 2., 3., 4., 0., 1., 0., 1., 2., 0., 1., 0., 1., 2., 3., 0., 1., 2., 3., 4.], axis=-1)) assert np.allclose(c0.buffer[:len(c0.buffer)].obs, obs) rew = [0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1] assert np.allclose(c0.buffer[:len(c0.buffer)].rew, [[x] * 4 for x in rew]) c2 = Collector(policy, envs, ReplayBuffer(size=100, stack_num=4), Logger.single_preprocess_fn, reward_metric=reward_metric) r = c2.collect(n_episode=[0, 0, 0, 10])['rew'] assert np.asanyarray(r).size == 1 and r == 4. batch = c2.sample(10) print(batch['obs_next']) if __name__ == '__main__': test_collector() test_collector_with_dict_state() test_collector_with_ma() test_collector_with_async()