from collections.abc import Callable, Sequence from test.base.env import MoveToRightEnv, NXEnv from typing import Any import gymnasium as gym import numpy as np import pytest import tqdm from tianshou.data import ( AsyncCollector, Batch, CachedReplayBuffer, Collector, PrioritizedReplayBuffer, ReplayBuffer, VectorReplayBuffer, ) from tianshou.data.batch import BatchProtocol from tianshou.data.types import ObsBatchProtocol, RolloutBatchProtocol from tianshou.env import DummyVectorEnv, SubprocVectorEnv from tianshou.policy import BasePolicy, TrainingStats try: import envpool except ImportError: envpool = None class MaxActionPolicy(BasePolicy): def __init__( self, action_space: gym.spaces.Space | None = None, dict_state: bool = False, need_state: bool = True, action_shape: Sequence[int] | int | None = None, ) -> None: """Mock policy for testing, will always return an array of ones of the shape of the action space. Note that this doesn't make much sense for discrete action space (the output is then intepreted as logits, meaning all actions would be equally likely). :param action_space: the action space of the environment. If None, a dummy Box space will be used. :param bool dict_state: if the observation of the environment is a dict :param bool need_state: if the policy needs the hidden state (for RNN) """ action_space = action_space or gym.spaces.Box(-1, 1, (1,)) super().__init__(action_space=action_space) self.dict_state = dict_state self.need_state = need_state self.action_shape = action_shape def forward( self, batch: ObsBatchProtocol, state: dict | BatchProtocol | np.ndarray | None = None, **kwargs: Any, ) -> Batch: if self.need_state: if state is None: state = np.zeros((len(batch.obs), 2)) elif isinstance(state, np.ndarray | BatchProtocol): state += np.int_(1) elif isinstance(state, dict) and state.get("hidden") is not None: state["hidden"] += np.int_(1) if self.dict_state: if self.action_shape: action_shape = self.action_shape elif isinstance(batch.obs, BatchProtocol): action_shape = len(batch.obs["index"]) else: action_shape = len(batch.obs) return Batch(act=np.ones(action_shape), state=state) action_shape = self.action_shape if self.action_shape else len(batch.obs) return Batch(act=np.ones(action_shape), state=state) def learn(self, batch: RolloutBatchProtocol, *args: Any, **kwargs: Any) -> TrainingStats: raise NotImplementedError def test_collector() -> None: env_fns = [lambda x=i: MoveToRightEnv(size=x, sleep=0) for i in [2, 3, 4, 5]] subproc_venv_4_envs = SubprocVectorEnv(env_fns) dummy_venv_4_envs = DummyVectorEnv(env_fns) policy = MaxActionPolicy() single_env = env_fns[0]() c_single_env = Collector( policy, single_env, ReplayBuffer(size=100), ) c_single_env.reset() c_single_env.collect(n_step=3) assert len(c_single_env.buffer) == 3 # TODO: direct attr access is an arcane way of using the buffer, it should be never done # The placeholders for entries are all zeros, so buffer.obs is an array filled with 3 # observations, and 97 zeros. # However, buffer[:] will have all attributes with length three... The non-filled entries are removed there # See above. For the single env, we start with obs=0, obs_next=1. # We move to obs=1, obs_next=2, # then the env is reset and we move to obs=0 # Making one more step results in obs_next=1 # The final 0 in the buffer.obs is because the buffer is initialized with zeros and the direct attr access assert np.allclose(c_single_env.buffer.obs[:4, 0], [0, 1, 0, 0]) obs_next = c_single_env.buffer[:].obs_next[..., 0] assert isinstance(obs_next, np.ndarray) assert np.allclose(obs_next, [1, 2, 1]) keys = np.zeros(100) keys[:3] = 1 assert np.allclose(c_single_env.buffer.info["key"], keys) for e in c_single_env.buffer.info["env"][:3]: assert isinstance(e, MoveToRightEnv) assert np.allclose(c_single_env.buffer.info["env_id"], 0) rews = np.zeros(100) rews[:3] = [0, 1, 0] assert np.allclose(c_single_env.buffer.rew, rews) # At this point, the buffer contains obs 0 -> 1 -> 0 # At start we have 3 entries in the buffer # We collect 3 episodes, in addition to the transitions we have collected before # 0 -> 1 -> 0 -> 0 (reset at collection start) -> 1 -> done (0) -> 1 -> done(0) # obs_next: 1 -> 2 -> 1 -> 1 (reset at collection start) -> 2 -> 1 -> 2 -> 1 -> 2 # In total, we will have 3 + 6 = 9 entries in the buffer c_single_env.collect(n_episode=3) assert len(c_single_env.buffer) == 8 assert np.allclose(c_single_env.buffer.obs[:10, 0], [0, 1, 0, 1, 0, 1, 0, 1, 0, 0]) obs_next = c_single_env.buffer[:].obs_next[..., 0] assert isinstance(obs_next, np.ndarray) assert np.allclose(obs_next, [1, 2, 1, 2, 1, 2, 1, 2]) assert np.allclose(c_single_env.buffer.info["key"][:8], 1) for e in c_single_env.buffer.info["env"][:8]: assert isinstance(e, MoveToRightEnv) assert np.allclose(c_single_env.buffer.info["env_id"][:8], 0) assert np.allclose(c_single_env.buffer.rew[:8], [0, 1, 0, 1, 0, 1, 0, 1]) c_single_env.collect(n_step=3, random=True) c_subproc_venv_4_envs = Collector( policy, subproc_venv_4_envs, VectorReplayBuffer(total_size=100, buffer_num=4), ) c_subproc_venv_4_envs.reset() # Collect some steps c_subproc_venv_4_envs.collect(n_step=8) obs = np.zeros(100) valid_indices = [0, 1, 25, 26, 50, 51, 75, 76] obs[valid_indices] = [0, 1, 0, 1, 0, 1, 0, 1] assert np.allclose(c_subproc_venv_4_envs.buffer.obs[:, 0], obs) obs_next = c_subproc_venv_4_envs.buffer[:].obs_next[..., 0] assert isinstance(obs_next, np.ndarray) assert np.allclose(obs_next, [1, 2, 1, 2, 1, 2, 1, 2]) keys = np.zeros(100) keys[valid_indices] = [1, 1, 1, 1, 1, 1, 1, 1] assert np.allclose(c_subproc_venv_4_envs.buffer.info["key"], keys) for e in c_subproc_venv_4_envs.buffer.info["env"][valid_indices]: assert isinstance(e, MoveToRightEnv) env_ids = np.zeros(100) env_ids[valid_indices] = [0, 0, 1, 1, 2, 2, 3, 3] assert np.allclose(c_subproc_venv_4_envs.buffer.info["env_id"], env_ids) rews = np.zeros(100) rews[valid_indices] = [0, 1, 0, 0, 0, 0, 0, 0] assert np.allclose(c_subproc_venv_4_envs.buffer.rew, rews) # we previously collected 8 steps, 2 from each env, now we collect 4 episodes # each env will contribute an episode, which will be of lens 2 (first env was reset), 1, 2, 3 # So we get 8 + 2+1+2+3 = 16 steps c_subproc_venv_4_envs.collect(n_episode=4) assert len(c_subproc_venv_4_envs.buffer) == 16 valid_indices = [2, 3, 27, 52, 53, 77, 78, 79] obs[valid_indices] = [0, 1, 2, 2, 3, 2, 3, 4] assert np.allclose(c_subproc_venv_4_envs.buffer.obs[:, 0], obs) obs_next = c_subproc_venv_4_envs.buffer[:].obs_next[..., 0] assert isinstance(obs_next, np.ndarray) assert np.allclose( obs_next, [1, 2, 1, 2, 1, 2, 3, 1, 2, 3, 4, 1, 2, 3, 4, 5], ) keys[valid_indices] = [1, 1, 1, 1, 1, 1, 1, 1] assert np.allclose(c_subproc_venv_4_envs.buffer.info["key"], keys) for e in c_subproc_venv_4_envs.buffer.info["env"][valid_indices]: assert isinstance(e, MoveToRightEnv) env_ids[valid_indices] = [0, 0, 1, 2, 2, 3, 3, 3] assert np.allclose(c_subproc_venv_4_envs.buffer.info["env_id"], env_ids) rews[valid_indices] = [0, 1, 1, 0, 1, 0, 0, 1] assert np.allclose(c_subproc_venv_4_envs.buffer.rew, rews) c_subproc_venv_4_envs.collect(n_episode=4, random=True) c_dummy_venv_4_envs = Collector( policy, dummy_venv_4_envs, VectorReplayBuffer(total_size=100, buffer_num=4), ) c_dummy_venv_4_envs.reset() c_dummy_venv_4_envs.collect(n_episode=7) obs1 = obs.copy() obs1[[4, 5, 28, 29, 30]] = [0, 1, 0, 1, 2] obs2 = obs.copy() obs2[[28, 29, 30, 54, 55, 56, 57]] = [0, 1, 2, 0, 1, 2, 3] c2obs = c_dummy_venv_4_envs.buffer.obs[:, 0] assert np.all(c2obs == obs1) or np.all(c2obs == obs2) c_dummy_venv_4_envs.reset_env() c_dummy_venv_4_envs.reset_buffer() assert c_dummy_venv_4_envs.collect(n_episode=8).n_collected_episodes == 8 valid_indices = [4, 5, 28, 29, 30, 54, 55, 56, 57] obs[valid_indices] = [0, 1, 0, 1, 2, 0, 1, 2, 3] assert np.all(c_dummy_venv_4_envs.buffer.obs[:, 0] == obs) keys[valid_indices] = [1, 1, 1, 1, 1, 1, 1, 1, 1] assert np.allclose(c_dummy_venv_4_envs.buffer.info["key"], keys) for e in c_dummy_venv_4_envs.buffer.info["env"][valid_indices]: assert isinstance(e, MoveToRightEnv) env_ids[valid_indices] = [0, 0, 1, 1, 1, 2, 2, 2, 2] assert np.allclose(c_dummy_venv_4_envs.buffer.info["env_id"], env_ids) rews[valid_indices] = [0, 1, 0, 0, 1, 0, 0, 0, 1] assert np.allclose(c_dummy_venv_4_envs.buffer.rew, rews) c_dummy_venv_4_envs.collect(n_episode=4, random=True) # test corner case with pytest.raises(ValueError): Collector(policy, dummy_venv_4_envs, ReplayBuffer(10)) with pytest.raises(ValueError): Collector(policy, dummy_venv_4_envs, PrioritizedReplayBuffer(10, 0.5, 0.5)) with pytest.raises(ValueError): c_dummy_venv_4_envs.collect() def get_env_factory(i: int, t: str) -> Callable[[], NXEnv]: return lambda: NXEnv(i, t) # test NXEnv for obs_type in ["array", "object"]: envs = SubprocVectorEnv([get_env_factory(i=i, t=obs_type) for i in [5, 10, 15, 20]]) c_suproc_new = Collector(policy, envs, VectorReplayBuffer(total_size=100, buffer_num=4)) c_suproc_new.reset() c_suproc_new.collect(n_step=6) assert c_suproc_new.buffer.obs.dtype == object @pytest.fixture() def async_collector_and_env_lens() -> tuple[AsyncCollector, list[int]]: env_lens = [2, 3, 4, 5] env_fns = [lambda x=i: MoveToRightEnv(size=x, sleep=0.001, random_sleep=True) for i in env_lens] venv = SubprocVectorEnv(env_fns, wait_num=len(env_fns) - 1) policy = MaxActionPolicy() bufsize = 60 async_collector = AsyncCollector( policy, venv, VectorReplayBuffer(total_size=bufsize * 4, buffer_num=4), ) async_collector.reset() return async_collector, env_lens class TestAsyncCollector: def test_collect_without_argument_gives_error( self, async_collector_and_env_lens: tuple[AsyncCollector, list[int]], ) -> None: c1, env_lens = async_collector_and_env_lens with pytest.raises(ValueError): c1.collect() def test_collect_one_episode_async( self, async_collector_and_env_lens: tuple[AsyncCollector, list[int]], ) -> None: c1, env_lens = async_collector_and_env_lens result = c1.collect(n_episode=1) assert result.n_collected_episodes >= 1 def test_enough_episodes_two_collection_cycles_n_episode_without_reset( self, async_collector_and_env_lens: tuple[AsyncCollector, list[int]], ) -> None: c1, env_lens = async_collector_and_env_lens n_episode = 2 result_c1 = c1.collect(n_episode=n_episode, reset_before_collect=False) assert result_c1.n_collected_episodes >= n_episode result_c2 = c1.collect(n_episode=n_episode, reset_before_collect=False) assert result_c2.n_collected_episodes >= n_episode def test_enough_episodes_two_collection_cycles_n_episode_with_reset( self, async_collector_and_env_lens: tuple[AsyncCollector, list[int]], ) -> None: c1, env_lens = async_collector_and_env_lens n_episode = 2 result_c1 = c1.collect(n_episode=n_episode, reset_before_collect=True) assert result_c1.n_collected_episodes >= n_episode result_c2 = c1.collect(n_episode=n_episode, reset_before_collect=True) assert result_c2.n_collected_episodes >= n_episode def test_enough_episodes_and_correct_obs_indices_and_obs_next_iterative_collection_cycles_n_episode( self, async_collector_and_env_lens: tuple[AsyncCollector, list[int]], ) -> None: c1, env_lens = async_collector_and_env_lens ptr = [0, 0, 0, 0] bufsize = 60 for n_episode in tqdm.trange(1, 30, desc="test async n_episode"): result = c1.collect(n_episode=n_episode) assert result.n_collected_episodes >= n_episode # check buffer data, obs and obs_next, env_id for i, count in enumerate(np.bincount(result.lens, minlength=6)[2:]): env_len = i + 2 total = env_len * count indices = np.arange(ptr[i], ptr[i] + total) % bufsize ptr[i] = (ptr[i] + total) % bufsize seq = np.arange(env_len) buf = c1.buffer.buffers[i] assert np.all(buf.info.env_id[indices] == i) assert np.all(buf.obs[indices].reshape(count, env_len) == seq) assert np.all(buf.obs_next[indices].reshape(count, env_len) == seq + 1) def test_enough_episodes_and_correct_obs_indices_and_obs_next_iterative_collection_cycles_n_step( self, async_collector_and_env_lens: tuple[AsyncCollector, list[int]], ) -> None: c1, env_lens = async_collector_and_env_lens bufsize = 60 ptr = [0, 0, 0, 0] for n_step in tqdm.trange(1, 15, desc="test async n_step"): result = c1.collect(n_step=n_step) assert result.n_collected_steps >= n_step for i, count in enumerate(np.bincount(result.lens, minlength=6)[2:]): env_len = i + 2 total = env_len * count indices = np.arange(ptr[i], ptr[i] + total) % bufsize ptr[i] = (ptr[i] + total) % bufsize seq = np.arange(env_len) buf = c1.buffer.buffers[i] assert np.all(buf.info.env_id[indices] == i) assert np.all(buf.obs[indices].reshape(count, env_len) == seq) assert np.all(buf.obs_next[indices].reshape(count, env_len) == seq + 1) def test_enough_episodes_and_correct_obs_indices_and_obs_next_iterative_collection_cycles_first_n_episode_then_n_step( self, async_collector_and_env_lens: tuple[AsyncCollector, list[int]], ) -> None: c1, env_lens = async_collector_and_env_lens bufsize = 60 ptr = [0, 0, 0, 0] for n_episode in tqdm.trange(1, 30, desc="test async n_episode"): result = c1.collect(n_episode=n_episode) assert result.n_collected_episodes >= n_episode # check buffer data, obs and obs_next, env_id for i, count in enumerate(np.bincount(result.lens, minlength=6)[2:]): env_len = i + 2 total = env_len * count indices = np.arange(ptr[i], ptr[i] + total) % bufsize ptr[i] = (ptr[i] + total) % bufsize seq = np.arange(env_len) buf = c1.buffer.buffers[i] assert np.all(buf.info.env_id[indices] == i) assert np.all(buf.obs[indices].reshape(count, env_len) == seq) assert np.all(buf.obs_next[indices].reshape(count, env_len) == seq + 1) # test async n_step, for now the buffer should be full of data, thus no bincount stuff as above for n_step in tqdm.trange(1, 15, desc="test async n_step"): result = c1.collect(n_step=n_step) assert result.n_collected_steps >= n_step for i in range(4): env_len = i + 2 seq = np.arange(env_len) buf = c1.buffer.buffers[i] assert np.all(buf.info.env_id == i) assert np.all(buf.obs.reshape(-1, env_len) == seq) assert np.all(buf.obs_next.reshape(-1, env_len) == seq + 1) def test_collector_with_dict_state() -> None: env = MoveToRightEnv(size=5, sleep=0, dict_state=True) policy = MaxActionPolicy(dict_state=True) c0 = Collector(policy, env, ReplayBuffer(size=100)) c0.reset() c0.collect(n_step=3) c0.collect(n_episode=2) assert len(c0.buffer) == 10 # 3 + two episodes with 5 steps each env_fns = [lambda x=i: MoveToRightEnv(size=x, sleep=0, dict_state=True) for i in [2, 3, 4, 5]] envs = DummyVectorEnv(env_fns) envs.seed(666) obs, info = envs.reset() assert not np.isclose(obs[0]["rand"], obs[1]["rand"]) c1 = Collector( policy, envs, VectorReplayBuffer(total_size=100, buffer_num=4), ) c1.reset() c1.collect(n_step=12) result = c1.collect(n_episode=8) assert result.n_collected_episodes == 8 lens = np.bincount(result.lens) assert ( result.n_collected_steps == 21 and np.all(lens == [0, 0, 2, 2, 2, 2]) or result.n_collected_steps == 20 and np.all(lens == [0, 0, 3, 1, 2, 2]) ) batch, _ = c1.buffer.sample(10) c0.buffer.update(c1.buffer) assert len(c0.buffer) in [42, 43] cur_obs = c0.buffer[:].obs assert isinstance(cur_obs, Batch) if len(c0.buffer) == 42: assert np.all( cur_obs.index[..., 0] == [ 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 2, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, ], ), cur_obs.index[..., 0] else: assert np.all( cur_obs.index[..., 0] == [ 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 0, 1, 0, 1, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, ], ), cur_obs.index[..., 0] c2 = Collector( policy, envs, VectorReplayBuffer(total_size=100, buffer_num=4, stack_num=4), ) c2.reset() c2.collect(n_episode=10) batch, _ = c2.buffer.sample(10) def test_collector_with_multi_agent() -> None: multi_agent_env = MoveToRightEnv(size=5, sleep=0, ma_rew=4) policy = MaxActionPolicy() c_single_env = Collector(policy, multi_agent_env, ReplayBuffer(size=100)) c_single_env.reset() multi_env_returns = c_single_env.collect(n_step=3).returns # c_single_env has length 3 # We have no full episodes, so no returns yet assert len(multi_env_returns) == 0 single_env_returns = c_single_env.collect(n_episode=2).returns # now two episodes. Since we have 4 a agents, the returns have shape (2, 4) assert single_env_returns.shape == (2, 4) assert np.all(single_env_returns == 1) env_fns = [lambda x=i: MoveToRightEnv(size=x, sleep=0, ma_rew=4) for i in [2, 3, 4, 5]] envs = DummyVectorEnv(env_fns) c_multi_env_ma = Collector( policy, envs, VectorReplayBuffer(total_size=100, buffer_num=4), ) c_multi_env_ma.reset() multi_env_returns = c_multi_env_ma.collect(n_step=12).returns # each env makes 3 steps, the first two envs are done and result in two finished episodes assert multi_env_returns.shape == (2, 4) and np.all(multi_env_returns == 1), multi_env_returns multi_env_returns = c_multi_env_ma.collect(n_episode=8).returns assert multi_env_returns.shape == (8, 4) assert np.all(multi_env_returns == 1) batch, _ = c_multi_env_ma.buffer.sample(10) print(batch) c_single_env.buffer.update(c_multi_env_ma.buffer) assert len(c_single_env.buffer) in [42, 43] if len(c_single_env.buffer) == 42: multi_env_returns = np.array( [ 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, ], ) else: multi_env_returns = np.array( [ 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, ], ) assert np.all(c_single_env.buffer[:].rew == [[x] * 4 for x in multi_env_returns]) assert np.all(c_single_env.buffer[:].done == multi_env_returns) c2 = Collector( policy, envs, VectorReplayBuffer(total_size=100, buffer_num=4, stack_num=4), ) c2.reset() multi_env_returns = c2.collect(n_episode=10).returns assert multi_env_returns.shape == (10, 4) assert np.all(multi_env_returns == 1) batch, _ = c2.buffer.sample(10) def test_collector_with_atari_setting() -> None: reference_obs = np.zeros([6, 4, 84, 84]) for i in range(6): reference_obs[i, 3, np.arange(84), np.arange(84)] = i reference_obs[i, 2, np.arange(84)] = i reference_obs[i, 1, :, np.arange(84)] = i reference_obs[i, 0] = i # atari single buffer env = MoveToRightEnv(size=5, sleep=0, array_state=True) policy = MaxActionPolicy() c0 = Collector(policy, env, ReplayBuffer(size=100)) c0.reset() c0.collect(n_step=6) c0.collect(n_episode=2) assert c0.buffer.obs.shape == (100, 4, 84, 84) assert c0.buffer.obs_next.shape == (100, 4, 84, 84) assert len(c0.buffer) == 15 # 6 + 2 episodes with 5 steps each obs = np.zeros_like(c0.buffer.obs) obs[np.arange(15)] = reference_obs[np.arange(15) % 5] assert np.all(obs == c0.buffer.obs) c1 = Collector(policy, env, ReplayBuffer(size=100, ignore_obs_next=True)) c1.collect(n_episode=3, reset_before_collect=True) assert np.allclose(c0.buffer.obs, c1.buffer.obs) with pytest.raises(AttributeError): c1.buffer.obs_next # noqa: B018 assert np.all(reference_obs[[1, 2, 3, 4, 4] * 3] == c1.buffer[:].obs_next) c2 = Collector( policy, env, ReplayBuffer(size=100, ignore_obs_next=True, save_only_last_obs=True), ) c2.reset() c2.collect(n_step=8) assert c2.buffer.obs.shape == (100, 84, 84) obs = np.zeros_like(c2.buffer.obs) obs[np.arange(8)] = reference_obs[[0, 1, 2, 3, 4, 0, 1, 2], -1] assert np.all(c2.buffer.obs == obs) obs_next = c2.buffer[:].obs_next assert isinstance(obs_next, np.ndarray) assert np.allclose(obs_next, reference_obs[[1, 2, 3, 4, 4, 1, 2, 2], -1]) # atari multi buffer env_fns = [lambda x=i: MoveToRightEnv(size=x, sleep=0, array_state=True) for i in [2, 3, 4, 5]] envs = DummyVectorEnv(env_fns) c3 = Collector(policy, envs, VectorReplayBuffer(total_size=100, buffer_num=4)) c3.reset() c3.collect(n_step=12) result = c3.collect(n_episode=9) assert result.n_collected_episodes == 9 assert result.n_collected_steps == 23 assert c3.buffer.obs.shape == (100, 4, 84, 84) obs = np.zeros_like(c3.buffer.obs) obs[np.arange(8)] = reference_obs[[0, 1, 0, 1, 0, 1, 0, 1]] obs[np.arange(25, 34)] = reference_obs[[0, 1, 2, 0, 1, 2, 0, 1, 2]] obs[np.arange(50, 58)] = reference_obs[[0, 1, 2, 3, 0, 1, 2, 3]] obs[np.arange(75, 85)] = reference_obs[[0, 1, 2, 3, 4, 0, 1, 2, 3, 4]] assert np.all(obs == c3.buffer.obs) obs_next = np.zeros_like(c3.buffer.obs_next) obs_next[np.arange(8)] = reference_obs[[1, 2, 1, 2, 1, 2, 1, 2]] obs_next[np.arange(25, 34)] = reference_obs[[1, 2, 3, 1, 2, 3, 1, 2, 3]] obs_next[np.arange(50, 58)] = reference_obs[[1, 2, 3, 4, 1, 2, 3, 4]] obs_next[np.arange(75, 85)] = reference_obs[[1, 2, 3, 4, 5, 1, 2, 3, 4, 5]] assert np.all(obs_next == c3.buffer.obs_next) c4 = Collector( policy, envs, VectorReplayBuffer( total_size=100, buffer_num=4, stack_num=4, ignore_obs_next=True, save_only_last_obs=True, ), ) c4.reset() c4.collect(n_step=12) result = c4.collect(n_episode=9) assert result.n_collected_episodes == 9 assert result.n_collected_steps == 23 assert c4.buffer.obs.shape == (100, 84, 84) obs = np.zeros_like(c4.buffer.obs) slice_obs = reference_obs[:, -1] obs[np.arange(8)] = slice_obs[[0, 1, 0, 1, 0, 1, 0, 1]] obs[np.arange(25, 34)] = slice_obs[[0, 1, 2, 0, 1, 2, 0, 1, 2]] obs[np.arange(50, 58)] = slice_obs[[0, 1, 2, 3, 0, 1, 2, 3]] obs[np.arange(75, 85)] = slice_obs[[0, 1, 2, 3, 4, 0, 1, 2, 3, 4]] assert np.all(c4.buffer.obs == obs) obs_next = np.zeros([len(c4.buffer), 4, 84, 84]) ref_index = np.array( [ 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 1, 2, 2, 1, 2, 2, 1, 2, 3, 3, 1, 2, 3, 3, 1, 2, 3, 4, 4, 1, 2, 3, 4, 4, ], ) obs_next[:, -1] = slice_obs[ref_index] ref_index -= 1 ref_index[ref_index < 0] = 0 obs_next[:, -2] = slice_obs[ref_index] ref_index -= 1 ref_index[ref_index < 0] = 0 obs_next[:, -3] = slice_obs[ref_index] ref_index -= 1 ref_index[ref_index < 0] = 0 obs_next[:, -4] = slice_obs[ref_index] assert np.all(obs_next == c4.buffer[:].obs_next) buf = ReplayBuffer(100, stack_num=4, ignore_obs_next=True, save_only_last_obs=True) c5 = Collector(policy, envs, CachedReplayBuffer(buf, 4, 10)) c5.reset() result_ = c5.collect(n_step=12) assert len(buf) == 5 assert len(c5.buffer) == 12 result = c5.collect(n_episode=9) assert result.n_collected_episodes == 9 assert result.n_collected_steps == 23 assert len(buf) == 35 assert np.all( buf.obs[: len(buf)] == slice_obs[ [ 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, 0, 1, 2, 3, 4, ] ], ) assert np.all( buf[:].obs_next[:, -1] == slice_obs[ [ 1, 1, 1, 2, 2, 1, 1, 1, 2, 3, 3, 1, 2, 3, 4, 4, 1, 1, 1, 2, 2, 1, 1, 1, 2, 3, 3, 1, 2, 2, 1, 2, 3, 4, 4, ] ], ) assert len(buf) == len(c5.buffer) # test buffer=None c6 = Collector(policy, envs) c6.reset() result1 = c6.collect(n_step=12) for key in ["n_collected_episodes", "n_collected_steps", "returns", "lens"]: assert np.allclose(getattr(result1, key), getattr(result_, key)) result2 = c6.collect(n_episode=9) for key in ["n_collected_episodes", "n_collected_steps", "returns", "lens"]: assert np.allclose(getattr(result2, key), getattr(result, key)) @pytest.mark.skipif(envpool is None, reason="EnvPool doesn't support this platform") def test_collector_envpool_gym_reset_return_info() -> None: envs = envpool.make_gymnasium("Pendulum-v1", num_envs=4, gym_reset_return_info=True) policy = MaxActionPolicy(action_shape=(len(envs), 1)) c0 = Collector( policy, envs, VectorReplayBuffer(len(envs) * 10, len(envs)), exploration_noise=True, ) c0.reset() c0.collect(n_step=8) env_ids = np.zeros(len(envs) * 10) env_ids[[0, 1, 10, 11, 20, 21, 30, 31]] = [0, 0, 1, 1, 2, 2, 3, 3] assert np.allclose(c0.buffer.info["env_id"], env_ids) def test_collector_with_vector_env() -> None: env_fns = [lambda x=i: MoveToRightEnv(size=x, sleep=0) for i in [1, 8, 9, 10]] dum = DummyVectorEnv(env_fns) policy = MaxActionPolicy() c2 = Collector( policy, dum, VectorReplayBuffer(total_size=100, buffer_num=4), ) c2.reset() c1r = c2.collect(n_episode=2) assert np.array_equal(np.array([1, 8]), c1r.lens) c2r = c2.collect(n_episode=10) assert np.array_equal(np.array([1, 1, 1, 1, 1, 1, 1, 8, 9, 10]), c2r.lens) c3r = c2.collect(n_step=20) assert np.array_equal(np.array([1, 1, 1, 1, 1]), c3r.lens) c4r = c2.collect(n_step=20) assert np.array_equal(np.array([1, 1, 1, 8, 1, 9, 1, 10]), c4r.lens) def test_async_collector_with_vector_env() -> None: env_fns = [lambda x=i: MoveToRightEnv(size=x, sleep=0) for i in [1, 8, 9, 10]] dum = DummyVectorEnv(env_fns) policy = MaxActionPolicy() c1 = AsyncCollector( policy, dum, VectorReplayBuffer(total_size=100, buffer_num=4), ) c1r = c1.collect(n_episode=10, reset_before_collect=True) assert np.array_equal(np.array([1, 1, 1, 1, 1, 1, 1, 1, 8, 1, 9]), c1r.lens) c2r = c1.collect(n_step=20) assert np.array_equal(np.array([1, 10, 1, 1, 1, 1]), c2r.lens)