07702fc007
# Goals of the PR
The PR introduces **no changes to functionality**, apart from improved
input validation here and there. The main goals are to reduce some
complexity of the code, to improve types and IDE completions, and to
extend documentation and block comments where appropriate. Because of
the change to the trainer interfaces, many files are affected (more
details below), but still the overall changes are "small" in a certain
sense.
## Major Change 1 - BatchProtocol
**TL;DR:** One can now annotate which fields the batch is expected to
have on input params and which fields a returned batch has. Should be
useful for reading the code. getting meaningful IDE support, and
catching bugs with mypy. This annotation strategy will continue to work
if Batch is replaced by TensorDict or by something else.
**In more detail:** Batch itself has no fields and using it for
annotations is of limited informational power. Batches with fields are
not separate classes but instead instances of Batch directly, so there
is no type that could be used for annotation. Fortunately, python
`Protocol` is here for the rescue. With these changes we can now do
things like
```python
class ActionBatchProtocol(BatchProtocol):
logits: Sequence[Union[tuple, torch.Tensor]]
dist: torch.distributions.Distribution
act: torch.Tensor
state: Optional[torch.Tensor]
class RolloutBatchProtocol(BatchProtocol):
obs: torch.Tensor
obs_next: torch.Tensor
info: Dict[str, Any]
rew: torch.Tensor
terminated: torch.Tensor
truncated: torch.Tensor
class PGPolicy(BasePolicy):
...
def forward(
self,
batch: RolloutBatchProtocol,
state: Optional[Union[dict, Batch, np.ndarray]] = None,
**kwargs: Any,
) -> ActionBatchProtocol:
```
The IDE and mypy are now very helpful in finding errors and in
auto-completion, whereas before the tools couldn't assist in that at
all.
## Major Change 2 - remove duplication in trainer package
**TL;DR:** There was a lot of duplication between `BaseTrainer` and its
subclasses. Even worse, it was almost-duplication. There was also
interface fragmentation through things like `onpolicy_trainer`. Now this
duplication is gone and all downstream code was adjusted.
**In more detail:** Since this change affects a lot of code, I would
like to explain why I thought it to be necessary.
1. The subclasses of `BaseTrainer` just duplicated docstrings and
constructors. What's worse, they changed the order of args there, even
turning some kwargs of BaseTrainer into args. They also had the arg
`learning_type` which was passed as kwarg to the base class and was
unused there. This made things difficult to maintain, and in fact some
errors were already present in the duplicated docstrings.
2. The "functions" a la `onpolicy_trainer`, which just called the
`OnpolicyTrainer.run`, not only introduced interface fragmentation but
also completely obfuscated the docstring and interfaces. They themselves
had no dosctring and the interface was just `*args, **kwargs`, which
makes it impossible to understand what they do and which things can be
passed without reading their implementation, then reading the docstring
of the associated class, etc. Needless to say, mypy and IDEs provide no
support with such functions. Nevertheless, they were used everywhere in
the code-base. I didn't find the sacrifices in clarity and complexity
justified just for the sake of not having to write `.run()` after
instantiating a trainer.
3. The trainers are all very similar to each other. As for my
application I needed a new trainer, I wanted to understand their
structure. The similarity, however, was hard to discover since they were
all in separate modules and there was so much duplication. I kept
staring at the constructors for a while until I figured out that
essentially no changes to the superclass were introduced. Now they are
all in the same module and the similarities/differences between them are
much easier to grasp (in my opinion)
4. Because of (1), I had to manually change and check a lot of code,
which was very tedious and boring. This kind of work won't be necessary
in the future, since now IDEs can be used for changing signatures,
renaming args and kwargs, changing class names and so on.
I have some more reasons, but maybe the above ones are convincing
enough.
## Minor changes: improved input validation and types
I added input validation for things like `state` and `action_scaling`
(which only makes sense for continuous envs). After adding this, some
tests failed to pass this validation. There I added
`action_scaling=isinstance(env.action_space, Box)`, after which tests
were green. I don't know why the tests were green before, since action
scaling doesn't make sense for discrete actions. I guess some aspect was
not tested and didn't crash.
I also added Literal in some places, in particular for
`action_bound_method`. Now it is no longer allowed to pass an empty
string, instead one should pass `None`. Also here there is input
validation with clear error messages.
@Trinkle23897 The functional tests are green. I didn't want to fix the
formatting, since it will change in the next PR that will solve #914
anyway. I also found a whole bunch of code in `docs/_static`, which I
just deleted (shouldn't it be copied from the sources during docs build
instead of committed?). I also haven't adjusted the documentation yet,
which atm still mentions the trainers of the type
`onpolicy_trainer(...)` instead of `OnpolicyTrainer(...).run()`
## Breaking Changes
The adjustments to the trainer package introduce breaking changes as
duplicated interfaces are deleted. However, it should be very easy for
users to adjust to them
---------
Co-authored-by: Michael Panchenko <m.panchenko@appliedai.de>
1407 lines
44 KiB
Python
1407 lines
44 KiB
Python
import os
|
|
import pickle
|
|
import tempfile
|
|
from timeit import timeit
|
|
|
|
import h5py
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from tianshou.data import (
|
|
Batch,
|
|
CachedReplayBuffer,
|
|
HERReplayBuffer,
|
|
HERVectorReplayBuffer,
|
|
PrioritizedReplayBuffer,
|
|
PrioritizedVectorReplayBuffer,
|
|
ReplayBuffer,
|
|
SegmentTree,
|
|
VectorReplayBuffer,
|
|
)
|
|
from tianshou.data.utils.converter import to_hdf5
|
|
|
|
if __name__ == '__main__':
|
|
from env import MyGoalEnv, MyTestEnv
|
|
else: # pytest
|
|
from test.base.env import MyGoalEnv, MyTestEnv
|
|
|
|
|
|
def test_replaybuffer(size=10, bufsize=20):
|
|
env = MyTestEnv(size)
|
|
buf = ReplayBuffer(bufsize)
|
|
buf.update(buf)
|
|
assert str(buf) == buf.__class__.__name__ + '()'
|
|
obs, _ = env.reset()
|
|
action_list = [1] * 5 + [0] * 10 + [1] * 10
|
|
for i, act in enumerate(action_list):
|
|
obs_next, rew, terminated, truncated, info = env.step(act)
|
|
buf.add(
|
|
Batch(
|
|
obs=obs,
|
|
act=[act],
|
|
rew=rew,
|
|
terminated=terminated,
|
|
truncated=truncated,
|
|
obs_next=obs_next,
|
|
info=info
|
|
)
|
|
)
|
|
obs = obs_next
|
|
assert len(buf) == min(bufsize, i + 1)
|
|
assert buf.act.dtype == int
|
|
assert buf.act.shape == (bufsize, 1)
|
|
data, indices = buf.sample(bufsize * 2)
|
|
assert (indices < len(buf)).all()
|
|
assert (data.obs < size).all()
|
|
assert (0 <= data.done).all() and (data.done <= 1).all()
|
|
assert (0 <= data.terminated).all() and (data.terminated <= 1).all()
|
|
assert (0 <= data.truncated).all() and (data.truncated <= 1).all()
|
|
b = ReplayBuffer(size=10)
|
|
# neg bsz should return empty index
|
|
assert b.sample_indices(-1).tolist() == []
|
|
ptr, ep_rew, ep_len, ep_idx = b.add(
|
|
Batch(
|
|
obs=1,
|
|
act=1,
|
|
rew=1,
|
|
terminated=1,
|
|
truncated=0,
|
|
obs_next='str',
|
|
info={
|
|
'a': 3,
|
|
'b': {
|
|
'c': 5.0
|
|
}
|
|
}
|
|
)
|
|
)
|
|
assert b.obs[0] == 1
|
|
assert b.done[0]
|
|
assert b.terminated[0]
|
|
assert not b.truncated[0]
|
|
assert b.obs_next[0] == 'str'
|
|
assert np.all(b.obs[1:] == 0)
|
|
assert np.all(b.obs_next[1:] == np.array(None))
|
|
assert b.info.a[0] == 3 and b.info.a.dtype == int
|
|
assert np.all(b.info.a[1:] == 0)
|
|
assert b.info.b.c[0] == 5.0 and b.info.b.c.dtype == float
|
|
assert np.all(b.info.b.c[1:] == 0.0)
|
|
assert ptr.shape == (1, ) and ptr[0] == 0
|
|
assert ep_rew.shape == (1, ) and ep_rew[0] == 1
|
|
assert ep_len.shape == (1, ) and ep_len[0] == 1
|
|
assert ep_idx.shape == (1, ) and ep_idx[0] == 0
|
|
# test extra keys pop up, the buffer should handle it dynamically
|
|
batch = Batch(
|
|
obs=2,
|
|
act=2,
|
|
rew=2,
|
|
terminated=0,
|
|
truncated=0,
|
|
obs_next="str2",
|
|
info={
|
|
"a": 4,
|
|
"d": {
|
|
"e": -np.inf
|
|
}
|
|
}
|
|
)
|
|
b.add(batch)
|
|
info_keys = ["a", "b", "d"]
|
|
assert set(b.info.keys()) == set(info_keys)
|
|
assert b.info.a[1] == 4 and b.info.b.c[1] == 0
|
|
assert b.info.d.e[1] == -np.inf
|
|
# test batch-style adding method, where len(batch) == 1
|
|
batch.done = [1]
|
|
batch.terminated = [0]
|
|
batch.truncated = [1]
|
|
batch.info.e = np.zeros([1, 4])
|
|
batch = Batch.stack([batch])
|
|
ptr, ep_rew, ep_len, ep_idx = b.add(batch, buffer_ids=[0])
|
|
assert ptr.shape == (1, ) and ptr[0] == 2
|
|
assert ep_rew.shape == (1, ) and ep_rew[0] == 4
|
|
assert ep_len.shape == (1, ) and ep_len[0] == 2
|
|
assert ep_idx.shape == (1, ) and ep_idx[0] == 1
|
|
assert set(b.info.keys()) == set(info_keys + ["e"])
|
|
assert b.info.e.shape == (b.maxsize, 1, 4)
|
|
with pytest.raises(IndexError):
|
|
b[22]
|
|
# test prev / next
|
|
assert np.all(b.prev(np.array([0, 1, 2])) == [0, 1, 1])
|
|
assert np.all(b.next(np.array([0, 1, 2])) == [0, 2, 2])
|
|
batch.done = [0]
|
|
b.add(batch, buffer_ids=[0])
|
|
assert np.all(b.prev(np.array([0, 1, 2, 3])) == [0, 1, 1, 3])
|
|
assert np.all(b.next(np.array([0, 1, 2, 3])) == [0, 2, 2, 3])
|
|
|
|
|
|
def test_ignore_obs_next(size=10):
|
|
# Issue 82
|
|
buf = ReplayBuffer(size, ignore_obs_next=True)
|
|
for i in range(size):
|
|
buf.add(
|
|
Batch(
|
|
obs={
|
|
'mask1': np.array([i, 1, 1, 0, 0]),
|
|
'mask2': np.array([i + 4, 0, 1, 0, 0]),
|
|
'mask': i
|
|
},
|
|
act={
|
|
'act_id': i,
|
|
'position_id': i + 3
|
|
},
|
|
rew=i,
|
|
terminated=i % 3 == 0,
|
|
truncated=False,
|
|
info={'if': i}
|
|
)
|
|
)
|
|
indices = np.arange(len(buf))
|
|
orig = np.arange(len(buf))
|
|
data = buf[indices]
|
|
data2 = buf[indices]
|
|
assert isinstance(data, Batch)
|
|
assert isinstance(data2, Batch)
|
|
assert np.allclose(indices, orig)
|
|
assert np.allclose(data.obs_next.mask, data2.obs_next.mask)
|
|
assert np.allclose(data.obs_next.mask, [0, 2, 3, 3, 5, 6, 6, 8, 9, 9])
|
|
buf.stack_num = 4
|
|
data = buf[indices]
|
|
data2 = buf[indices]
|
|
assert np.allclose(data.obs_next.mask, data2.obs_next.mask)
|
|
assert np.allclose(
|
|
data.obs_next.mask,
|
|
np.array(
|
|
[
|
|
[0, 0, 0, 0], [1, 1, 1, 2], [1, 1, 2, 3], [1, 1, 2, 3], [4, 4, 4, 5],
|
|
[4, 4, 5, 6], [4, 4, 5, 6], [7, 7, 7, 8], [7, 7, 8, 9], [7, 7, 8, 9]
|
|
]
|
|
)
|
|
)
|
|
assert np.allclose(data.info['if'], data2.info['if'])
|
|
assert np.allclose(
|
|
data.info['if'],
|
|
np.array(
|
|
[
|
|
[0, 0, 0, 0], [1, 1, 1, 1], [1, 1, 1, 2], [1, 1, 2, 3], [4, 4, 4, 4],
|
|
[4, 4, 4, 5], [4, 4, 5, 6], [7, 7, 7, 7], [7, 7, 7, 8], [7, 7, 8, 9]
|
|
]
|
|
)
|
|
)
|
|
assert data.obs_next
|
|
|
|
|
|
def test_stack(size=5, bufsize=9, stack_num=4, cached_num=3):
|
|
env = MyTestEnv(size)
|
|
buf = ReplayBuffer(bufsize, stack_num=stack_num)
|
|
buf2 = ReplayBuffer(bufsize, stack_num=stack_num, sample_avail=True)
|
|
buf3 = ReplayBuffer(bufsize, stack_num=stack_num, save_only_last_obs=True)
|
|
obs, info = env.reset(options={"state": 1})
|
|
for _ in range(16):
|
|
obs_next, rew, terminated, truncated, info = env.step(1)
|
|
done = terminated or truncated
|
|
buf.add(
|
|
Batch(
|
|
obs=obs,
|
|
act=1,
|
|
rew=rew,
|
|
terminated=terminated,
|
|
truncated=truncated,
|
|
info=info
|
|
)
|
|
)
|
|
buf2.add(
|
|
Batch(
|
|
obs=obs,
|
|
act=1,
|
|
rew=rew,
|
|
terminated=terminated,
|
|
truncated=truncated,
|
|
info=info
|
|
)
|
|
)
|
|
buf3.add(
|
|
Batch(
|
|
obs=[obs, obs, obs],
|
|
act=1,
|
|
rew=rew,
|
|
terminated=terminated,
|
|
truncated=truncated,
|
|
obs_next=[obs, obs],
|
|
info=info
|
|
)
|
|
)
|
|
obs = obs_next
|
|
if done:
|
|
obs, info = env.reset(options={"state": 1})
|
|
indices = np.arange(len(buf))
|
|
assert np.allclose(
|
|
buf.get(indices, 'obs')[..., 0], [
|
|
[1, 1, 1, 2], [1, 1, 2, 3], [1, 2, 3, 4], [1, 1, 1, 1], [1, 1, 1, 2],
|
|
[1, 1, 2, 3], [1, 2, 3, 4], [4, 4, 4, 4], [1, 1, 1, 1]
|
|
]
|
|
)
|
|
assert np.allclose(buf.get(indices, 'obs'), buf3.get(indices, 'obs'))
|
|
assert np.allclose(buf.get(indices, 'obs'), buf3.get(indices, 'obs_next'))
|
|
_, indices = buf2.sample(0)
|
|
assert indices.tolist() == [2, 6]
|
|
_, indices = buf2.sample(1)
|
|
assert indices[0] in [2, 6]
|
|
batch, indices = buf2.sample(-1) # neg bsz -> no data
|
|
assert indices.tolist() == [] and len(batch) == 0
|
|
with pytest.raises(IndexError):
|
|
buf[bufsize * 2]
|
|
|
|
|
|
def test_priortized_replaybuffer(size=32, bufsize=15):
|
|
env = MyTestEnv(size)
|
|
buf = PrioritizedReplayBuffer(bufsize, 0.5, 0.5)
|
|
buf2 = PrioritizedVectorReplayBuffer(bufsize, buffer_num=3, alpha=0.5, beta=0.5)
|
|
obs, info = env.reset()
|
|
action_list = [1] * 5 + [0] * 10 + [1] * 10
|
|
for i, act in enumerate(action_list):
|
|
obs_next, rew, terminated, truncated, info = env.step(act)
|
|
batch = Batch(
|
|
obs=obs,
|
|
act=act,
|
|
rew=rew,
|
|
terminated=terminated,
|
|
truncated=truncated,
|
|
obs_next=obs_next,
|
|
info=info,
|
|
policy=np.random.randn() - 0.5
|
|
)
|
|
batch_stack = Batch.stack([batch, batch, batch])
|
|
buf.add(Batch.stack([batch]), buffer_ids=[0])
|
|
buf2.add(batch_stack, buffer_ids=[0, 1, 2])
|
|
obs = obs_next
|
|
data, indices = buf.sample(len(buf) // 2)
|
|
if len(buf) // 2 == 0:
|
|
assert len(data) == len(buf)
|
|
else:
|
|
assert len(data) == len(buf) // 2
|
|
assert len(buf) == min(bufsize, i + 1)
|
|
assert len(buf2) == min(bufsize, 3 * (i + 1))
|
|
# check single buffer's data
|
|
assert buf.info.key.shape == (buf.maxsize, )
|
|
assert buf.rew.dtype == float
|
|
assert buf.done.dtype == bool
|
|
assert buf.terminated.dtype == bool
|
|
assert buf.truncated.dtype == bool
|
|
data, indices = buf.sample(len(buf) // 2)
|
|
buf.update_weight(indices, -data.weight / 2)
|
|
assert np.allclose(buf.weight[indices], np.abs(-data.weight / 2)**buf._alpha)
|
|
# check multi buffer's data
|
|
assert np.allclose(buf2[np.arange(buf2.maxsize)].weight, 1)
|
|
batch, indices = buf2.sample(10)
|
|
buf2.update_weight(indices, batch.weight * 0)
|
|
weight = buf2[np.arange(buf2.maxsize)].weight
|
|
mask = np.isin(np.arange(buf2.maxsize), indices)
|
|
assert np.all(weight[mask] == weight[mask][0])
|
|
assert np.all(weight[~mask] == weight[~mask][0])
|
|
assert weight[~mask][0] < weight[mask][0] and weight[mask][0] <= 1
|
|
|
|
|
|
def test_herreplaybuffer(size=10, bufsize=100, sample_sz=4):
|
|
env_size = size
|
|
env = MyGoalEnv(env_size, array_state=True)
|
|
|
|
def compute_reward_fn(ag, g):
|
|
return env.compute_reward_fn(ag, g, {})
|
|
|
|
buf = HERReplayBuffer(
|
|
bufsize, compute_reward_fn=compute_reward_fn, horizon=30, future_k=8
|
|
)
|
|
buf2 = HERVectorReplayBuffer(
|
|
bufsize,
|
|
buffer_num=3,
|
|
compute_reward_fn=compute_reward_fn,
|
|
horizon=30,
|
|
future_k=8
|
|
)
|
|
# Apply her on every episodes sampled (Hacky but necessary for deterministic test)
|
|
buf.future_p = 1
|
|
for buf2_buf in buf2.buffers:
|
|
buf2_buf.future_p = 1
|
|
|
|
obs, _ = env.reset()
|
|
action_list = [1] * 5 + [0] * 10 + [1] * 10
|
|
for i, act in enumerate(action_list):
|
|
obs_next, rew, terminated, truncated, info = env.step(act)
|
|
batch = Batch(
|
|
obs=obs,
|
|
act=[act],
|
|
rew=rew,
|
|
terminated=terminated,
|
|
truncated=truncated,
|
|
obs_next=obs_next,
|
|
info=info
|
|
)
|
|
buf.add(batch)
|
|
buf2.add(Batch.stack([batch, batch, batch]), buffer_ids=[0, 1, 2])
|
|
obs = obs_next
|
|
assert len(buf) == min(bufsize, i + 1)
|
|
assert len(buf2) == min(bufsize, 3 * (i + 1))
|
|
|
|
batch, indices = buf.sample(sample_sz)
|
|
|
|
# Check that goals are the same for the episode (only 1 ep in buffer)
|
|
tmp_indices = indices.copy()
|
|
for _ in range(2 * env_size):
|
|
obs = buf[tmp_indices].obs
|
|
obs_next = buf[tmp_indices].obs_next
|
|
rew = buf[tmp_indices].rew
|
|
g = obs.desired_goal.reshape(sample_sz, -1)[:, 0]
|
|
ag_next = obs_next.achieved_goal.reshape(sample_sz, -1)[:, 0]
|
|
g_next = obs_next.desired_goal.reshape(sample_sz, -1)[:, 0]
|
|
assert np.all(g == g[0])
|
|
assert np.all(g_next == g_next[0])
|
|
assert np.all(rew == (ag_next == g).astype(np.float32))
|
|
tmp_indices = buf.next(tmp_indices)
|
|
|
|
# Check that goals are correctly restored
|
|
buf._restore_cache()
|
|
tmp_indices = indices.copy()
|
|
for _ in range(2 * env_size):
|
|
obs = buf[tmp_indices].obs
|
|
obs_next = buf[tmp_indices].obs_next
|
|
g = obs.desired_goal.reshape(sample_sz, -1)[:, 0]
|
|
g_next = obs_next.desired_goal.reshape(sample_sz, -1)[:, 0]
|
|
assert np.all(g == env_size)
|
|
assert np.all(g_next == g_next[0])
|
|
assert np.all(g == g[0])
|
|
tmp_indices = buf.next(tmp_indices)
|
|
|
|
# Test vector buffer
|
|
batch, indices = buf2.sample(sample_sz)
|
|
|
|
# Check that goals are the same for the episode (only 1 ep in buffer)
|
|
tmp_indices = indices.copy()
|
|
for _ in range(2 * env_size):
|
|
obs = buf2[tmp_indices].obs
|
|
obs_next = buf2[tmp_indices].obs_next
|
|
rew = buf2[tmp_indices].rew
|
|
g = obs.desired_goal.reshape(sample_sz, -1)[:, 0]
|
|
ag_next = obs_next.achieved_goal.reshape(sample_sz, -1)[:, 0]
|
|
g_next = obs_next.desired_goal.reshape(sample_sz, -1)[:, 0]
|
|
assert np.all(g == g_next)
|
|
assert np.all(rew == (ag_next == g).astype(np.float32))
|
|
tmp_indices = buf2.next(tmp_indices)
|
|
|
|
# Check that goals are correctly restored
|
|
buf2._restore_cache()
|
|
tmp_indices = indices.copy()
|
|
for _ in range(2 * env_size):
|
|
obs = buf2[tmp_indices].obs
|
|
obs_next = buf2[tmp_indices].obs_next
|
|
g = obs.desired_goal.reshape(sample_sz, -1)[:, 0]
|
|
g_next = obs_next.desired_goal.reshape(sample_sz, -1)[:, 0]
|
|
assert np.all(g == env_size)
|
|
assert np.all(g_next == g_next[0])
|
|
assert np.all(g == g[0])
|
|
tmp_indices = buf2.next(tmp_indices)
|
|
|
|
# Test handling cycled indices
|
|
env_size = size
|
|
bufsize = 15
|
|
env = MyGoalEnv(env_size, array_state=False)
|
|
|
|
def compute_reward_fn(ag, g):
|
|
return env.compute_reward_fn(ag, g, {})
|
|
|
|
buf = HERReplayBuffer(
|
|
bufsize, compute_reward_fn=compute_reward_fn, horizon=30, future_k=8
|
|
)
|
|
buf._index = 5 # shifted start index
|
|
buf.future_p = 1
|
|
action_list = [1] * 10
|
|
for ep_len in [5, 10]:
|
|
obs, _ = env.reset()
|
|
for i in range(ep_len):
|
|
act = 1
|
|
obs_next, rew, terminated, truncated, info = env.step(act)
|
|
batch = Batch(
|
|
obs=obs,
|
|
act=[act],
|
|
rew=rew,
|
|
terminated=(i == ep_len - 1),
|
|
truncated=(i == ep_len - 1),
|
|
obs_next=obs_next,
|
|
info=info
|
|
)
|
|
buf.add(batch)
|
|
obs = obs_next
|
|
batch, indices = buf.sample(0)
|
|
assert np.all(buf[:5].obs.desired_goal == buf[0].obs.desired_goal)
|
|
assert np.all(buf[5:10].obs.desired_goal == buf[5].obs.desired_goal)
|
|
assert np.all(buf[10:].obs.desired_goal == buf[0].obs.desired_goal) # (same ep)
|
|
assert np.all(buf[0].obs.desired_goal != buf[5].obs.desired_goal) # (diff ep)
|
|
|
|
# Another test case for cycled indices
|
|
env_size = 99
|
|
bufsize = 15
|
|
env = MyGoalEnv(env_size, array_state=False)
|
|
buf = HERReplayBuffer(
|
|
bufsize, compute_reward_fn=compute_reward_fn, horizon=30, future_k=8
|
|
)
|
|
buf.future_p = 1
|
|
for x, ep_len in enumerate([10, 20]):
|
|
obs, _ = env.reset()
|
|
for i in range(ep_len):
|
|
act = 1
|
|
obs_next, rew, terminated, truncated, info = env.step(act)
|
|
batch = Batch(
|
|
obs=obs,
|
|
act=[act],
|
|
rew=rew,
|
|
terminated=(i == ep_len - 1),
|
|
truncated=(i == ep_len - 1),
|
|
obs_next=obs_next,
|
|
info=info
|
|
)
|
|
if x == 1 and obs["observation"] < 10:
|
|
obs = obs_next
|
|
continue
|
|
buf.add(batch)
|
|
obs = obs_next
|
|
buf._restore_cache()
|
|
sample_indices = np.array([10]) # Suppose the sampled indices is [10]
|
|
buf.rewrite_transitions(sample_indices)
|
|
assert int(buf.obs.desired_goal[10][0]) in [11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
|
|
|
|
|
|
def test_update():
|
|
buf1 = ReplayBuffer(4, stack_num=2)
|
|
buf2 = ReplayBuffer(4, stack_num=2)
|
|
for i in range(5):
|
|
buf1.add(
|
|
Batch(
|
|
obs=np.array([i]),
|
|
act=float(i),
|
|
rew=i * i,
|
|
terminated=i % 2 == 0,
|
|
truncated=False,
|
|
info={'incident': 'found'}
|
|
)
|
|
)
|
|
assert len(buf1) > len(buf2)
|
|
buf2.update(buf1)
|
|
assert len(buf1) == len(buf2)
|
|
assert (buf2[0].obs == buf1[1].obs).all()
|
|
assert (buf2[-1].obs == buf1[0].obs).all()
|
|
b = CachedReplayBuffer(ReplayBuffer(10), 4, 5)
|
|
with pytest.raises(NotImplementedError):
|
|
b.update(b)
|
|
|
|
|
|
def test_segtree():
|
|
realop = np.sum
|
|
# small test
|
|
actual_len = 8
|
|
tree = SegmentTree(actual_len) # 1-15. 8-15 are leaf nodes
|
|
assert len(tree) == actual_len
|
|
assert np.all([tree[i] == 0. for i in range(actual_len)])
|
|
with pytest.raises(IndexError):
|
|
tree[actual_len]
|
|
naive = np.zeros([actual_len])
|
|
for _ in range(1000):
|
|
# random choose a place to perform single update
|
|
index = np.random.randint(actual_len)
|
|
value = np.random.rand()
|
|
naive[index] = value
|
|
tree[index] = value
|
|
for i in range(actual_len):
|
|
for j in range(i + 1, actual_len):
|
|
ref = realop(naive[i:j])
|
|
out = tree.reduce(i, j)
|
|
assert np.allclose(ref, out), (ref, out)
|
|
assert np.allclose(tree.reduce(start=1), realop(naive[1:]))
|
|
assert np.allclose(tree.reduce(end=-1), realop(naive[:-1]))
|
|
# batch setitem
|
|
for _ in range(1000):
|
|
index = np.random.choice(actual_len, size=4)
|
|
value = np.random.rand(4)
|
|
naive[index] = value
|
|
tree[index] = value
|
|
assert np.allclose(realop(naive), tree.reduce())
|
|
for _ in range(10):
|
|
left = np.random.randint(actual_len)
|
|
right = np.random.randint(left + 1, actual_len + 1)
|
|
assert np.allclose(realop(naive[left:right]), tree.reduce(left, right))
|
|
# large test
|
|
actual_len = 16384
|
|
tree = SegmentTree(actual_len)
|
|
naive = np.zeros([actual_len])
|
|
for _ in range(1000):
|
|
index = np.random.choice(actual_len, size=64)
|
|
value = np.random.rand(64)
|
|
naive[index] = value
|
|
tree[index] = value
|
|
assert np.allclose(realop(naive), tree.reduce())
|
|
for _ in range(10):
|
|
left = np.random.randint(actual_len)
|
|
right = np.random.randint(left + 1, actual_len + 1)
|
|
assert np.allclose(realop(naive[left:right]), tree.reduce(left, right))
|
|
|
|
# test prefix-sum-idx
|
|
actual_len = 8
|
|
tree = SegmentTree(actual_len)
|
|
naive = np.random.rand(actual_len)
|
|
tree[np.arange(actual_len)] = naive
|
|
for _ in range(1000):
|
|
scalar = np.random.rand() * naive.sum()
|
|
index = tree.get_prefix_sum_idx(scalar)
|
|
assert naive[:index].sum() <= scalar <= naive[:index + 1].sum()
|
|
# corner case here
|
|
naive = np.ones(actual_len, int)
|
|
tree[np.arange(actual_len)] = naive
|
|
for scalar in range(actual_len):
|
|
index = tree.get_prefix_sum_idx(scalar * 1.)
|
|
assert naive[:index].sum() <= scalar <= naive[:index + 1].sum()
|
|
tree = SegmentTree(10)
|
|
tree[np.arange(3)] = np.array([0.1, 0, 0.1])
|
|
assert np.allclose(
|
|
tree.get_prefix_sum_idx(np.array([0, .1, .1 + 1e-6, .2 - 1e-6])), [0, 0, 2, 2]
|
|
)
|
|
with pytest.raises(AssertionError):
|
|
tree.get_prefix_sum_idx(.2)
|
|
# test large prefix-sum-idx
|
|
actual_len = 16384
|
|
tree = SegmentTree(actual_len)
|
|
naive = np.random.rand(actual_len)
|
|
tree[np.arange(actual_len)] = naive
|
|
for _ in range(1000):
|
|
scalar = np.random.rand() * naive.sum()
|
|
index = tree.get_prefix_sum_idx(scalar)
|
|
assert naive[:index].sum() <= scalar <= naive[:index + 1].sum()
|
|
|
|
# profile
|
|
if __name__ == '__main__':
|
|
size = 100000
|
|
bsz = 64
|
|
naive = np.random.rand(size)
|
|
tree = SegmentTree(size)
|
|
tree[np.arange(size)] = naive
|
|
|
|
def sample_npbuf():
|
|
return np.random.choice(size, bsz, p=naive / naive.sum())
|
|
|
|
def sample_tree():
|
|
scalar = np.random.rand(bsz) * tree.reduce()
|
|
return tree.get_prefix_sum_idx(scalar)
|
|
|
|
print('npbuf', timeit(sample_npbuf, setup=sample_npbuf, number=1000))
|
|
print('tree', timeit(sample_tree, setup=sample_tree, number=1000))
|
|
|
|
|
|
def test_pickle():
|
|
size = 100
|
|
vbuf = ReplayBuffer(size, stack_num=2)
|
|
pbuf = PrioritizedReplayBuffer(size, 0.6, 0.4)
|
|
rew = np.array([1, 1])
|
|
for i in range(4):
|
|
vbuf.add(
|
|
Batch(
|
|
obs=Batch(index=np.array([i])),
|
|
act=0,
|
|
rew=rew,
|
|
terminated=0,
|
|
truncated=0,
|
|
)
|
|
)
|
|
for i in range(5):
|
|
pbuf.add(
|
|
Batch(
|
|
obs=Batch(index=np.array([i])),
|
|
act=2,
|
|
rew=rew,
|
|
terminated=0,
|
|
truncated=0,
|
|
info=np.random.rand()
|
|
)
|
|
)
|
|
# save & load
|
|
_vbuf = pickle.loads(pickle.dumps(vbuf))
|
|
_pbuf = pickle.loads(pickle.dumps(pbuf))
|
|
assert len(_vbuf) == len(vbuf) and np.allclose(_vbuf.act, vbuf.act)
|
|
assert len(_pbuf) == len(pbuf) and np.allclose(_pbuf.act, pbuf.act)
|
|
# make sure the meta var is identical
|
|
assert _vbuf.stack_num == vbuf.stack_num
|
|
assert np.allclose(
|
|
_pbuf.weight[np.arange(len(_pbuf))], pbuf.weight[np.arange(len(pbuf))]
|
|
)
|
|
|
|
|
|
def test_hdf5():
|
|
size = 100
|
|
buffers = {
|
|
"array": ReplayBuffer(size, stack_num=2),
|
|
"prioritized": PrioritizedReplayBuffer(size, 0.6, 0.4),
|
|
}
|
|
buffer_types = {k: b.__class__ for k, b in buffers.items()}
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
info_t = torch.tensor([1.]).to(device)
|
|
for i in range(4):
|
|
kwargs = {
|
|
'obs': Batch(index=np.array([i])),
|
|
'act': i,
|
|
'rew': np.array([1, 2]),
|
|
'terminated': i % 3 == 2,
|
|
'truncated': False,
|
|
'done': i % 3 == 2,
|
|
'info': {
|
|
"number": {
|
|
"n": i,
|
|
"t": info_t
|
|
},
|
|
'extra': None
|
|
},
|
|
}
|
|
buffers["array"].add(Batch(kwargs))
|
|
buffers["prioritized"].add(Batch(kwargs))
|
|
|
|
# save
|
|
paths = {}
|
|
for k, buf in buffers.items():
|
|
f, path = tempfile.mkstemp(suffix='.hdf5')
|
|
os.close(f)
|
|
buf.save_hdf5(path)
|
|
paths[k] = path
|
|
|
|
# load replay buffer
|
|
_buffers = {k: buffer_types[k].load_hdf5(paths[k]) for k in paths.keys()}
|
|
|
|
# compare
|
|
for k in buffers.keys():
|
|
assert len(_buffers[k]) == len(buffers[k])
|
|
assert np.allclose(_buffers[k].act, buffers[k].act)
|
|
assert _buffers[k].stack_num == buffers[k].stack_num
|
|
assert _buffers[k].maxsize == buffers[k].maxsize
|
|
assert np.all(_buffers[k]._indices == buffers[k]._indices)
|
|
for k in ["array", "prioritized"]:
|
|
assert _buffers[k]._index == buffers[k]._index
|
|
assert isinstance(buffers[k].get(0, "info"), Batch)
|
|
assert isinstance(_buffers[k].get(0, "info"), Batch)
|
|
for k in ["array"]:
|
|
assert np.all(buffers[k][:].info.number.n == _buffers[k][:].info.number.n)
|
|
assert np.all(buffers[k][:].info.extra == _buffers[k][:].info.extra)
|
|
|
|
# raise exception when value cannot be pickled
|
|
data = {"not_supported": lambda x: x * x}
|
|
grp = h5py.Group
|
|
with pytest.raises(NotImplementedError):
|
|
to_hdf5(data, grp)
|
|
# ndarray with data type not supported by HDF5 that cannot be pickled
|
|
data = {"not_supported": np.array(lambda x: x * x)}
|
|
grp = h5py.Group
|
|
with pytest.raises(RuntimeError):
|
|
to_hdf5(data, grp)
|
|
|
|
|
|
def test_replaybuffermanager():
|
|
buf = VectorReplayBuffer(20, 4)
|
|
batch = Batch(
|
|
obs=[1, 2, 3],
|
|
act=[1, 2, 3],
|
|
rew=[1, 2, 3],
|
|
terminated=[0, 0, 1],
|
|
truncated=[0, 0, 0],
|
|
)
|
|
ptr, ep_rew, ep_len, ep_idx = buf.add(batch, buffer_ids=[0, 1, 2])
|
|
assert np.all(ep_len == [0, 0, 1]) and np.all(ep_rew == [0, 0, 3])
|
|
assert np.all(ptr == [0, 5, 10]) and np.all(ep_idx == [0, 5, 10])
|
|
with pytest.raises(NotImplementedError):
|
|
# ReplayBufferManager cannot be updated
|
|
buf.update(buf)
|
|
# sample index / prev / next / unfinished_index
|
|
indices = buf.sample_indices(11000)
|
|
assert np.bincount(indices)[[0, 5, 10]].min() >= 3000 # uniform sample
|
|
batch, indices = buf.sample(0)
|
|
assert np.allclose(indices, [0, 5, 10])
|
|
indices_prev = buf.prev(indices)
|
|
assert np.allclose(indices_prev, indices), indices_prev
|
|
indices_next = buf.next(indices)
|
|
assert np.allclose(indices_next, indices), indices_next
|
|
assert np.allclose(buf.unfinished_index(), [0, 5])
|
|
buf.add(
|
|
Batch(obs=[4], act=[4], rew=[4], terminated=[1], truncated=[0]), buffer_ids=[3]
|
|
)
|
|
assert np.allclose(buf.unfinished_index(), [0, 5])
|
|
batch, indices = buf.sample(10)
|
|
batch, indices = buf.sample(0)
|
|
assert np.allclose(indices, [0, 5, 10, 15])
|
|
indices_prev = buf.prev(indices)
|
|
assert np.allclose(indices_prev, indices), indices_prev
|
|
indices_next = buf.next(indices)
|
|
assert np.allclose(indices_next, indices), indices_next
|
|
data = np.array([0, 0, 0, 0])
|
|
buf.add(
|
|
Batch(obs=data, act=data, rew=data, terminated=data, truncated=data),
|
|
buffer_ids=[0, 1, 2, 3]
|
|
)
|
|
buf.add(
|
|
Batch(obs=data, act=data, rew=data, terminated=1 - data, truncated=data),
|
|
buffer_ids=[0, 1, 2, 3]
|
|
)
|
|
assert len(buf) == 12
|
|
buf.add(
|
|
Batch(obs=data, act=data, rew=data, terminated=data, truncated=data),
|
|
buffer_ids=[0, 1, 2, 3]
|
|
)
|
|
buf.add(
|
|
Batch(obs=data, act=data, rew=data, terminated=[0, 1, 0, 1], truncated=data),
|
|
buffer_ids=[0, 1, 2, 3]
|
|
)
|
|
assert len(buf) == 20
|
|
indices = buf.sample_indices(120000)
|
|
assert np.bincount(indices).min() >= 5000
|
|
batch, indices = buf.sample(10)
|
|
indices = buf.sample_indices(0)
|
|
assert np.allclose(indices, np.arange(len(buf)))
|
|
# check the actual data stored in buf._meta
|
|
assert np.allclose(
|
|
buf.done, [
|
|
0,
|
|
0,
|
|
1,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
1,
|
|
0,
|
|
1,
|
|
1,
|
|
0,
|
|
1,
|
|
0,
|
|
0,
|
|
1,
|
|
0,
|
|
1,
|
|
0,
|
|
1,
|
|
]
|
|
)
|
|
assert np.allclose(
|
|
buf.prev(indices), [
|
|
0,
|
|
0,
|
|
1,
|
|
3,
|
|
3,
|
|
5,
|
|
5,
|
|
6,
|
|
8,
|
|
8,
|
|
10,
|
|
11,
|
|
11,
|
|
13,
|
|
13,
|
|
15,
|
|
16,
|
|
16,
|
|
18,
|
|
18,
|
|
]
|
|
)
|
|
assert np.allclose(
|
|
buf.next(indices), [
|
|
1,
|
|
2,
|
|
2,
|
|
4,
|
|
4,
|
|
6,
|
|
7,
|
|
7,
|
|
9,
|
|
9,
|
|
10,
|
|
12,
|
|
12,
|
|
14,
|
|
14,
|
|
15,
|
|
17,
|
|
17,
|
|
19,
|
|
19,
|
|
]
|
|
)
|
|
assert np.allclose(buf.unfinished_index(), [4, 14])
|
|
ptr, ep_rew, ep_len, ep_idx = buf.add(
|
|
Batch(obs=[1], act=[1], rew=[1], terminated=[1], truncated=[0]), buffer_ids=[2]
|
|
)
|
|
assert np.all(ep_len == [3]) and np.all(ep_rew == [1])
|
|
assert np.all(ptr == [10]) and np.all(ep_idx == [13])
|
|
assert np.allclose(buf.unfinished_index(), [4])
|
|
indices = list(sorted(buf.sample_indices(0)))
|
|
assert np.allclose(indices, np.arange(len(buf)))
|
|
assert np.allclose(
|
|
buf.prev(indices), [
|
|
0,
|
|
0,
|
|
1,
|
|
3,
|
|
3,
|
|
5,
|
|
5,
|
|
6,
|
|
8,
|
|
8,
|
|
14,
|
|
11,
|
|
11,
|
|
13,
|
|
13,
|
|
15,
|
|
16,
|
|
16,
|
|
18,
|
|
18,
|
|
]
|
|
)
|
|
assert np.allclose(
|
|
buf.next(indices), [
|
|
1,
|
|
2,
|
|
2,
|
|
4,
|
|
4,
|
|
6,
|
|
7,
|
|
7,
|
|
9,
|
|
9,
|
|
10,
|
|
12,
|
|
12,
|
|
14,
|
|
10,
|
|
15,
|
|
17,
|
|
17,
|
|
19,
|
|
19,
|
|
]
|
|
)
|
|
# corner case: list, int and -1
|
|
assert buf.prev(-1) == buf.prev([buf.maxsize - 1])[0]
|
|
assert buf.next(-1) == buf.next([buf.maxsize - 1])[0]
|
|
batch = buf._meta
|
|
batch.info = np.ones(buf.maxsize)
|
|
buf.set_batch(batch)
|
|
assert np.allclose(buf.buffers[-1].info, [1] * 5)
|
|
assert buf.sample_indices(-1).tolist() == []
|
|
assert np.array([ReplayBuffer(0, ignore_obs_next=True)]).dtype == object
|
|
|
|
|
|
def test_cachedbuffer():
|
|
buf = CachedReplayBuffer(ReplayBuffer(10), 4, 5)
|
|
assert buf.sample_indices(0).tolist() == []
|
|
# check the normal function/usage/storage in CachedReplayBuffer
|
|
ptr, ep_rew, ep_len, ep_idx = buf.add(
|
|
Batch(obs=[1], act=[1], rew=[1], terminated=[0], truncated=[0]), buffer_ids=[1]
|
|
)
|
|
obs = np.zeros(buf.maxsize)
|
|
obs[15] = 1
|
|
indices = buf.sample_indices(0)
|
|
assert np.allclose(indices, [15])
|
|
assert np.allclose(buf.prev(indices), [15])
|
|
assert np.allclose(buf.next(indices), [15])
|
|
assert np.allclose(buf.obs, obs)
|
|
assert np.all(ep_len == [0]) and np.all(ep_rew == [0.0])
|
|
assert np.all(ptr == [15]) and np.all(ep_idx == [15])
|
|
ptr, ep_rew, ep_len, ep_idx = buf.add(
|
|
Batch(obs=[2], act=[2], rew=[2], terminated=[1], truncated=[0]), buffer_ids=[3]
|
|
)
|
|
obs[[0, 25]] = 2
|
|
indices = buf.sample_indices(0)
|
|
assert np.allclose(indices, [0, 15])
|
|
assert np.allclose(buf.prev(indices), [0, 15])
|
|
assert np.allclose(buf.next(indices), [0, 15])
|
|
assert np.allclose(buf.obs, obs)
|
|
assert np.all(ep_len == [1]) and np.all(ep_rew == [2.0])
|
|
assert np.all(ptr == [0]) and np.all(ep_idx == [0])
|
|
assert np.allclose(buf.unfinished_index(), [15])
|
|
assert np.allclose(buf.sample_indices(0), [0, 15])
|
|
ptr, ep_rew, ep_len, ep_idx = buf.add(
|
|
Batch(obs=[3, 4], act=[3, 4], rew=[3, 4], terminated=[0, 1], truncated=[0, 0]),
|
|
buffer_ids=[3, 1] # TODO
|
|
)
|
|
assert np.all(ep_len == [0, 2]) and np.all(ep_rew == [0, 5.0])
|
|
assert np.all(ptr == [25, 2]) and np.all(ep_idx == [25, 1])
|
|
obs[[0, 1, 2, 15, 16, 25]] = [2, 1, 4, 1, 4, 3]
|
|
assert np.allclose(buf.obs, obs)
|
|
assert np.allclose(buf.unfinished_index(), [25])
|
|
indices = buf.sample_indices(0)
|
|
assert np.allclose(indices, [0, 1, 2, 25])
|
|
assert np.allclose(buf.done[indices], [1, 0, 1, 0])
|
|
assert np.allclose(buf.prev(indices), [0, 1, 1, 25])
|
|
assert np.allclose(buf.next(indices), [0, 2, 2, 25])
|
|
indices = buf.sample_indices(10000)
|
|
assert np.bincount(indices)[[0, 1, 2, 25]].min() > 2000 # uniform sample
|
|
# cached buffer with main_buffer size == 0 (no update)
|
|
# used in test_collector
|
|
buf = CachedReplayBuffer(ReplayBuffer(0, sample_avail=True), 4, 5)
|
|
data = np.zeros(4)
|
|
rew = np.ones([4, 4])
|
|
buf.add(
|
|
Batch(
|
|
obs=data,
|
|
act=data,
|
|
rew=rew,
|
|
terminated=[0, 0, 1, 1],
|
|
truncated=[0, 0, 0, 0]
|
|
)
|
|
)
|
|
buf.add(
|
|
Batch(
|
|
obs=data,
|
|
act=data,
|
|
rew=rew,
|
|
terminated=[0, 0, 0, 0],
|
|
truncated=[0, 0, 0, 0]
|
|
)
|
|
)
|
|
buf.add(
|
|
Batch(
|
|
obs=data,
|
|
act=data,
|
|
rew=rew,
|
|
terminated=[1, 1, 1, 1],
|
|
truncated=[0, 0, 0, 0]
|
|
)
|
|
)
|
|
buf.add(
|
|
Batch(
|
|
obs=data,
|
|
act=data,
|
|
rew=rew,
|
|
terminated=[0, 0, 0, 0],
|
|
truncated=[0, 0, 0, 0]
|
|
)
|
|
)
|
|
ptr, ep_rew, ep_len, ep_idx = buf.add(
|
|
Batch(
|
|
obs=data,
|
|
act=data,
|
|
rew=rew,
|
|
terminated=[0, 1, 0, 1],
|
|
truncated=[0, 0, 0, 0]
|
|
)
|
|
)
|
|
assert np.all(ptr == [1, -1, 11, -1]) and np.all(ep_idx == [0, -1, 10, -1])
|
|
assert np.all(ep_len == [0, 2, 0, 2])
|
|
assert np.all(ep_rew == [data, data + 2, data, data + 2])
|
|
assert np.allclose(
|
|
buf.done, [
|
|
0,
|
|
0,
|
|
1,
|
|
0,
|
|
0,
|
|
0,
|
|
1,
|
|
1,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
1,
|
|
0,
|
|
0,
|
|
0,
|
|
]
|
|
)
|
|
indices = buf.sample_indices(0)
|
|
assert np.allclose(indices, [0, 1, 10, 11])
|
|
assert np.allclose(buf.prev(indices), [0, 0, 10, 10])
|
|
assert np.allclose(buf.next(indices), [1, 1, 11, 11])
|
|
|
|
|
|
def test_multibuf_stack():
|
|
size = 5
|
|
bufsize = 9
|
|
stack_num = 4
|
|
cached_num = 3
|
|
env = MyTestEnv(size)
|
|
# test if CachedReplayBuffer can handle stack_num + ignore_obs_next
|
|
buf4 = CachedReplayBuffer(
|
|
ReplayBuffer(bufsize, stack_num=stack_num, ignore_obs_next=True), cached_num,
|
|
size
|
|
)
|
|
# test if CachedReplayBuffer can handle corner case:
|
|
# buffer + stack_num + ignore_obs_next + sample_avail
|
|
buf5 = CachedReplayBuffer(
|
|
ReplayBuffer(
|
|
bufsize, stack_num=stack_num, ignore_obs_next=True, sample_avail=True
|
|
), cached_num, size
|
|
)
|
|
obs, info = env.reset(options={"state": 1})
|
|
for i in range(18):
|
|
obs_next, rew, terminated, truncated, info = env.step(1)
|
|
done = terminated or truncated
|
|
obs_list = np.array([obs + size * i for i in range(cached_num)])
|
|
act_list = [1] * cached_num
|
|
rew_list = [rew] * cached_num
|
|
terminated_list = [terminated] * cached_num
|
|
truncated_list = [truncated] * cached_num
|
|
obs_next_list = -obs_list
|
|
info_list = [info] * cached_num
|
|
batch = Batch(
|
|
obs=obs_list,
|
|
act=act_list,
|
|
rew=rew_list,
|
|
terminated=terminated_list,
|
|
truncated=truncated_list,
|
|
obs_next=obs_next_list,
|
|
info=info_list
|
|
)
|
|
buf5.add(batch)
|
|
buf4.add(batch)
|
|
assert np.all(buf4.obs == buf5.obs)
|
|
assert np.all(buf4.done == buf5.done)
|
|
assert np.all(buf4.terminated == buf5.terminated)
|
|
assert np.all(buf4.truncated == buf5.truncated)
|
|
obs = obs_next
|
|
if done:
|
|
obs, info = env.reset(options={"state": 1})
|
|
# check the `add` order is correct
|
|
assert np.allclose(
|
|
buf4.obs.reshape(-1),
|
|
[
|
|
12,
|
|
13,
|
|
14,
|
|
4,
|
|
6,
|
|
7,
|
|
8,
|
|
9,
|
|
11, # main_buffer
|
|
1,
|
|
2,
|
|
3,
|
|
4,
|
|
0, # cached_buffer[0]
|
|
6,
|
|
7,
|
|
8,
|
|
9,
|
|
0, # cached_buffer[1]
|
|
11,
|
|
12,
|
|
13,
|
|
14,
|
|
0, # cached_buffer[2]
|
|
]
|
|
), buf4.obs
|
|
assert np.allclose(
|
|
buf4.done,
|
|
[
|
|
0,
|
|
0,
|
|
1,
|
|
1,
|
|
0,
|
|
0,
|
|
0,
|
|
1,
|
|
0, # main_buffer
|
|
0,
|
|
0,
|
|
0,
|
|
1,
|
|
0, # cached_buffer[0]
|
|
0,
|
|
0,
|
|
0,
|
|
1,
|
|
0, # cached_buffer[1]
|
|
0,
|
|
0,
|
|
0,
|
|
1,
|
|
0, # cached_buffer[2]
|
|
]
|
|
), buf4.done
|
|
assert np.allclose(buf4.unfinished_index(), [10, 15, 20])
|
|
indices = sorted(buf4.sample_indices(0))
|
|
assert np.allclose(indices, list(range(bufsize)) + [9, 10, 14, 15, 19, 20])
|
|
assert np.allclose(
|
|
buf4[indices].obs[..., 0], [
|
|
[11, 11, 11, 12],
|
|
[11, 11, 12, 13],
|
|
[11, 12, 13, 14],
|
|
[4, 4, 4, 4],
|
|
[6, 6, 6, 6],
|
|
[6, 6, 6, 7],
|
|
[6, 6, 7, 8],
|
|
[6, 7, 8, 9],
|
|
[11, 11, 11, 11],
|
|
[1, 1, 1, 1],
|
|
[1, 1, 1, 2],
|
|
[6, 6, 6, 6],
|
|
[6, 6, 6, 7],
|
|
[11, 11, 11, 11],
|
|
[11, 11, 11, 12],
|
|
]
|
|
)
|
|
assert np.allclose(
|
|
buf4[indices].obs_next[..., 0], [
|
|
[11, 11, 12, 13],
|
|
[11, 12, 13, 14],
|
|
[11, 12, 13, 14],
|
|
[4, 4, 4, 4],
|
|
[6, 6, 6, 7],
|
|
[6, 6, 7, 8],
|
|
[6, 7, 8, 9],
|
|
[6, 7, 8, 9],
|
|
[11, 11, 11, 12],
|
|
[1, 1, 1, 2],
|
|
[1, 1, 1, 2],
|
|
[6, 6, 6, 7],
|
|
[6, 6, 6, 7],
|
|
[11, 11, 11, 12],
|
|
[11, 11, 11, 12],
|
|
]
|
|
)
|
|
indices = buf5.sample_indices(0)
|
|
assert np.allclose(sorted(indices), [2, 7])
|
|
assert np.all(np.isin(buf5.sample_indices(100), indices))
|
|
# manually change the stack num
|
|
buf5.stack_num = 2
|
|
for buf in buf5.buffers:
|
|
buf.stack_num = 2
|
|
indices = buf5.sample_indices(0)
|
|
assert np.allclose(sorted(indices), [0, 1, 2, 5, 6, 7, 10, 15, 20])
|
|
batch, _ = buf5.sample(0)
|
|
# test Atari with CachedReplayBuffer, save_only_last_obs + ignore_obs_next
|
|
buf6 = CachedReplayBuffer(
|
|
ReplayBuffer(
|
|
bufsize, stack_num=stack_num, save_only_last_obs=True, ignore_obs_next=True
|
|
), cached_num, size
|
|
)
|
|
obs = np.random.rand(size, 4, 84, 84)
|
|
buf6.add(
|
|
Batch(
|
|
obs=[obs[2], obs[0]],
|
|
act=[1, 1],
|
|
rew=[0, 0],
|
|
terminated=[0, 1],
|
|
truncated=[0, 0],
|
|
obs_next=[obs[3], obs[1]]
|
|
),
|
|
buffer_ids=[1, 2]
|
|
)
|
|
assert buf6.obs.shape == (buf6.maxsize, 84, 84)
|
|
assert np.allclose(buf6.obs[0], obs[0, -1])
|
|
assert np.allclose(buf6.obs[14], obs[2, -1])
|
|
assert np.allclose(buf6.obs[19], obs[0, -1])
|
|
assert buf6[0].obs.shape == (4, 84, 84)
|
|
|
|
|
|
def test_multibuf_hdf5():
|
|
size = 100
|
|
buffers = {
|
|
"vector": VectorReplayBuffer(size * 4, 4),
|
|
"cached": CachedReplayBuffer(ReplayBuffer(size), 4, size)
|
|
}
|
|
buffer_types = {k: b.__class__ for k, b in buffers.items()}
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
info_t = torch.tensor([1.]).to(device)
|
|
for i in range(4):
|
|
kwargs = {
|
|
'obs': Batch(index=np.array([i])),
|
|
'act': i,
|
|
'rew': np.array([1, 2]),
|
|
'terminated': i % 3 == 2,
|
|
'truncated': False,
|
|
'done': i % 3 == 2,
|
|
'info': {
|
|
"number": {
|
|
"n": i,
|
|
"t": info_t
|
|
},
|
|
'extra': None
|
|
},
|
|
}
|
|
buffers["vector"].add(
|
|
Batch.stack([kwargs, kwargs, kwargs]), buffer_ids=[0, 1, 2]
|
|
)
|
|
buffers["cached"].add(
|
|
Batch.stack([kwargs, kwargs, kwargs]), buffer_ids=[0, 1, 2]
|
|
)
|
|
|
|
# save
|
|
paths = {}
|
|
for k, buf in buffers.items():
|
|
f, path = tempfile.mkstemp(suffix='.hdf5')
|
|
os.close(f)
|
|
buf.save_hdf5(path)
|
|
paths[k] = path
|
|
|
|
# load replay buffer
|
|
_buffers = {k: buffer_types[k].load_hdf5(paths[k]) for k in paths.keys()}
|
|
|
|
# compare
|
|
for k in buffers.keys():
|
|
assert len(_buffers[k]) == len(buffers[k])
|
|
assert np.allclose(_buffers[k].act, buffers[k].act)
|
|
assert _buffers[k].stack_num == buffers[k].stack_num
|
|
assert _buffers[k].maxsize == buffers[k].maxsize
|
|
assert np.all(_buffers[k]._indices == buffers[k]._indices)
|
|
# check shallow copy in VectorReplayBuffer
|
|
for k in ["vector", "cached"]:
|
|
buffers[k].info.number.n[0] = -100
|
|
assert buffers[k].buffers[0].info.number.n[0] == -100
|
|
# check if still behave normally
|
|
for k in ["vector", "cached"]:
|
|
kwargs = {
|
|
'obs': Batch(index=np.array([5])),
|
|
'act': 5,
|
|
'rew': np.array([2, 1]),
|
|
'terminated': False,
|
|
'truncated': False,
|
|
'done': False,
|
|
'info': {
|
|
"number": {
|
|
"n": i
|
|
},
|
|
'Timelimit.truncate': True
|
|
},
|
|
}
|
|
buffers[k].add(Batch.stack([kwargs, kwargs, kwargs, kwargs]))
|
|
act = np.zeros(buffers[k].maxsize)
|
|
if k == "vector":
|
|
act[np.arange(5)] = np.array([0, 1, 2, 3, 5])
|
|
act[np.arange(5) + size] = np.array([0, 1, 2, 3, 5])
|
|
act[np.arange(5) + size * 2] = np.array([0, 1, 2, 3, 5])
|
|
act[size * 3] = 5
|
|
elif k == "cached":
|
|
act[np.arange(9)] = np.array([0, 1, 2, 0, 1, 2, 0, 1, 2])
|
|
act[np.arange(3) + size] = np.array([3, 5, 2])
|
|
act[np.arange(3) + size * 2] = np.array([3, 5, 2])
|
|
act[np.arange(3) + size * 3] = np.array([3, 5, 2])
|
|
act[size * 4] = 5
|
|
assert np.allclose(buffers[k].act, act)
|
|
info_keys = ["number", "extra", "Timelimit.truncate"]
|
|
assert set(buffers[k].info.keys()) == set(info_keys)
|
|
|
|
for path in paths.values():
|
|
os.remove(path)
|
|
|
|
|
|
def test_from_data():
|
|
obs_data = np.ndarray((10, 3, 3), dtype="uint8")
|
|
for i in range(10):
|
|
obs_data[i] = i * np.ones((3, 3), dtype="uint8")
|
|
obs_next_data = np.zeros_like(obs_data)
|
|
obs_next_data[:-1] = obs_data[1:]
|
|
f, path = tempfile.mkstemp(suffix='.hdf5')
|
|
os.close(f)
|
|
with h5py.File(path, "w") as f:
|
|
obs = f.create_dataset("obs", data=obs_data)
|
|
act = f.create_dataset("act", data=np.arange(10, dtype="int32"))
|
|
rew = f.create_dataset("rew", data=np.arange(10, dtype="float32"))
|
|
terminated = f.create_dataset("terminated", data=np.zeros(10, dtype="bool"))
|
|
truncated = f.create_dataset("truncated", data=np.zeros(10, dtype="bool"))
|
|
done = f.create_dataset("done", data=np.zeros(10, dtype="bool"))
|
|
obs_next = f.create_dataset("obs_next", data=obs_next_data)
|
|
buf = ReplayBuffer.from_data(
|
|
obs, act, rew, terminated, truncated, done, obs_next
|
|
)
|
|
assert len(buf) == 10
|
|
batch = buf[3]
|
|
assert np.array_equal(batch.obs, 3 * np.ones((3, 3), dtype="uint8"))
|
|
assert batch.act == 3
|
|
assert batch.rew == 3.0
|
|
assert not batch.done
|
|
assert np.array_equal(batch.obs_next, 4 * np.ones((3, 3), dtype="uint8"))
|
|
os.remove(path)
|
|
|
|
|
|
def test_custom_key():
|
|
batch = Batch(
|
|
**{
|
|
'obs_next':
|
|
np.array(
|
|
[
|
|
[
|
|
1.174, -0.1151, -0.609, -0.5205, -0.9316, 3.236, -2.418, 0.386,
|
|
0.2227, -0.5117, 2.293
|
|
]
|
|
]
|
|
),
|
|
'rew':
|
|
np.array([4.28125]),
|
|
'act':
|
|
np.array([[-0.3088, -0.4636, 0.4956]]),
|
|
'truncated':
|
|
np.array([False]),
|
|
'obs':
|
|
np.array(
|
|
[
|
|
[
|
|
1.193, -0.1203, -0.6123, -0.519, -0.9434, 3.32, -2.266, 0.9116,
|
|
0.623, 0.1259, 0.363
|
|
]
|
|
]
|
|
),
|
|
'terminated':
|
|
np.array([False]),
|
|
'done':
|
|
np.array([False]),
|
|
'returns':
|
|
np.array([74.70343082]),
|
|
'info':
|
|
Batch(),
|
|
'policy':
|
|
Batch(),
|
|
}
|
|
)
|
|
buffer_size = len(batch.rew)
|
|
buffer = ReplayBuffer(buffer_size)
|
|
buffer.add(batch)
|
|
sampled_batch, _ = buffer.sample(1)
|
|
# Check if they have the same keys
|
|
assert set(batch.keys()) == set(sampled_batch.keys()), \
|
|
"Batches have different keys: {} and {}".format(
|
|
set(batch.keys()), set(sampled_batch.keys()))
|
|
# Compare the values for each key
|
|
for key in batch.keys():
|
|
if isinstance(batch.__dict__[key], np.ndarray
|
|
) and isinstance(sampled_batch.__dict__[key], np.ndarray):
|
|
assert np.allclose(batch.__dict__[key], sampled_batch.__dict__[key]), \
|
|
"Value mismatch for key: {}".format(key)
|
|
if isinstance(batch.__dict__[key],
|
|
Batch) and isinstance(sampled_batch.__dict__[key], Batch):
|
|
assert batch.__dict__[key].is_empty()
|
|
assert sampled_batch.__dict__[key].is_empty()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_replaybuffer()
|
|
test_ignore_obs_next()
|
|
test_stack()
|
|
test_segtree()
|
|
test_priortized_replaybuffer()
|
|
test_update()
|
|
test_pickle()
|
|
test_hdf5()
|
|
test_replaybuffermanager()
|
|
test_cachedbuffer()
|
|
test_multibuf_stack()
|
|
test_multibuf_hdf5()
|
|
test_from_data()
|
|
test_herreplaybuffer()
|
|
test_custom_key()
|