wehub-resource-sync
189 lines
7.5 KiB
Python
189 lines
7.5 KiB
Python
import unittest
|
|
|
|
import numpy as np
|
|
|
|
from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch
|
|
from ray.rllib.utils.framework import try_import_tf
|
|
from ray.rllib.utils.minibatch_utils import (
|
|
MiniBatchCyclicIterator,
|
|
ShardEpisodesIterator,
|
|
)
|
|
from ray.rllib.utils.test_utils import check
|
|
from ray.rllib.utils.torch_utils import convert_to_torch_tensor
|
|
|
|
tf1, tf, tfv = try_import_tf()
|
|
tf1.enable_eager_execution()
|
|
|
|
CONFIGS = [
|
|
{"minibatch_size": 256, "num_epochs": 30, "agent_steps": (1652, 1463)},
|
|
{"minibatch_size": 128, "num_epochs": 10, "agent_steps": (1000, 2)},
|
|
{"minibatch_size": 128, "num_epochs": 3, "agent_steps": (56, 56)},
|
|
{"minibatch_size": 128, "num_epochs": 7, "agent_steps": (56, 56)},
|
|
{"minibatch_size": 128, "num_epochs": 10, "agent_steps": (56, 56)},
|
|
{"minibatch_size": 128, "num_epochs": 10, "agent_steps": (56, 3)},
|
|
{"minibatch_size": 128, "num_epochs": 10, "agent_steps": (56, 4)},
|
|
{"minibatch_size": 128, "num_epochs": 10, "agent_steps": (56, 55)},
|
|
{"minibatch_size": 128, "num_epochs": 10, "agent_steps": (400, 400)},
|
|
{"minibatch_size": 128, "num_epochs": 10, "agent_steps": (64, 64)},
|
|
# W/ SEQ_LENS.
|
|
{
|
|
"minibatch_size": 64,
|
|
"num_epochs": 1,
|
|
"agent_steps": (128,),
|
|
"seq_lens": [16, 16, 16, 16, 16, 16, 2, 2, 14, 14],
|
|
"padding": True,
|
|
},
|
|
]
|
|
|
|
|
|
class TestMinibatchUtils(unittest.TestCase):
|
|
def test_minibatch_cyclic_iterator(self):
|
|
|
|
for config in CONFIGS:
|
|
minibatch_size = config["minibatch_size"]
|
|
num_epochs = config["num_epochs"]
|
|
agent_steps = config["agent_steps"]
|
|
seq_lens = config.get("seq_lens")
|
|
max_seq_len = None
|
|
if seq_lens:
|
|
max_seq_len = max(seq_lens)
|
|
padding = config.get("padding", False)
|
|
num_env_steps = max(agent_steps)
|
|
|
|
for backend in ["torch", "numpy"]:
|
|
sample_batches = {
|
|
f"pol{i}": SampleBatch(
|
|
{
|
|
"obs": np.arange(agent_steps[i]),
|
|
"seq_lens": seq_lens,
|
|
}
|
|
)
|
|
if not seq_lens or not padding
|
|
else SampleBatch(
|
|
{
|
|
"obs": np.concatenate(
|
|
[
|
|
np.concatenate(
|
|
[
|
|
np.arange(s),
|
|
np.zeros(shape=(max_seq_len - s,)),
|
|
]
|
|
)
|
|
for s in seq_lens
|
|
]
|
|
),
|
|
"seq_lens": seq_lens,
|
|
},
|
|
_zero_padded=padding,
|
|
)
|
|
for i in range(len(agent_steps))
|
|
}
|
|
if backend == "torch":
|
|
for pid, batch in sample_batches.items():
|
|
batch["obs"] = convert_to_torch_tensor(batch["obs"])
|
|
if seq_lens:
|
|
batch["seq_lens"] = convert_to_torch_tensor(
|
|
batch["seq_lens"]
|
|
)
|
|
|
|
mb = MultiAgentBatch(sample_batches, num_env_steps)
|
|
batch_iter = MiniBatchCyclicIterator(
|
|
mb,
|
|
minibatch_size=minibatch_size,
|
|
num_epochs=num_epochs,
|
|
shuffle_batch_per_epoch=False,
|
|
)
|
|
print(config)
|
|
iteration_counter = 0
|
|
for batch in batch_iter:
|
|
print(batch)
|
|
print("-" * 80)
|
|
print(batch["pol0"]["obs"])
|
|
print("*" * 80)
|
|
# Check that for each policy the batch size is equal to the
|
|
# minibatch_size.
|
|
for policy_batch in batch.policy_batches.values():
|
|
check(policy_batch.count, minibatch_size)
|
|
iteration_counter += 1
|
|
|
|
# For each policy check that the last item in batch matches the expected
|
|
# values, i.e. iteration_counter * minibatch_size % agent_steps - 1.
|
|
total_steps = iteration_counter * minibatch_size
|
|
for policy_idx, policy_batch in enumerate(
|
|
batch.policy_batches.values()
|
|
):
|
|
expected_last_item = (total_steps - 1) % agent_steps[policy_idx]
|
|
if seq_lens and seq_lens[-1] < max_seq_len:
|
|
expected_last_item = 0.0
|
|
check(policy_batch["obs"][-1], expected_last_item)
|
|
|
|
# Check iteration counter (should be
|
|
# ceil(num_gsd_iter * max(agent_steps) / minibatch_size)).
|
|
expected_iteration_counter = np.ceil(
|
|
num_epochs * max(agent_steps) / minibatch_size
|
|
)
|
|
if not seq_lens:
|
|
check(iteration_counter, expected_iteration_counter)
|
|
print(f"iteration_counter: {iteration_counter}")
|
|
|
|
def test_shard_episodes_iterator(self):
|
|
class DummyEpisode:
|
|
def __init__(self, length):
|
|
self.length = length
|
|
# Dummy data to represent the episode content.
|
|
self.data = [0] * length
|
|
|
|
def __len__(self):
|
|
return self.length
|
|
|
|
def __getitem__(self, key):
|
|
assert isinstance(key, slice)
|
|
return self.slice(key)
|
|
|
|
def slice(self, slice, len_lookback_buffer=None):
|
|
# Create a new Episode object with the sliced length
|
|
return DummyEpisode(len(self.data[slice]))
|
|
|
|
def __repr__(self):
|
|
return f"{(type(self).__name__)}({self.length})"
|
|
|
|
# Create a list of episodes with varying lengths
|
|
episode_lens = [10, 21, 3, 4, 35, 41, 5, 15, 44]
|
|
|
|
episodes = [DummyEpisode(len_) for len_ in episode_lens]
|
|
|
|
# Number of shards
|
|
num_shards = 3
|
|
# Create the iterator
|
|
iterator = ShardEpisodesIterator(episodes, num_shards)
|
|
# Iterate and collect the results
|
|
shards = list(iterator)
|
|
# The sharder should try to split as few times as possible. In our
|
|
# case here, only the len=4 episode is split into 1 and 3. All other
|
|
# episodes are kept as-is. Yet, the resulting sub-lists have all
|
|
# either size 59 or 60.
|
|
check([len(e) for e in shards[0]], [44, 10, 5]) # 59
|
|
check([len(e) for e in shards[1]], [41, 15, 3]) # 59
|
|
check([len(e) for e in shards[2]], [35, 21, 1, 3]) # 60
|
|
|
|
# Different number of shards.
|
|
num_shards = 4
|
|
# Create the iterator.
|
|
iterator = ShardEpisodesIterator(episodes, num_shards)
|
|
# Iterate and collect the results
|
|
shards = list(iterator)
|
|
# The sharder should try to split as few times as possible, keeping
|
|
# as many episodes as-is (w/o splitting).
|
|
check([len(e) for e in shards[0]], [44]) # 44
|
|
check([len(e) for e in shards[1]], [41, 3]) # 44
|
|
check([len(e) for e in shards[2]], [35, 10]) # 45
|
|
check([len(e) for e in shards[3]], [21, 15, 5, 1, 3]) # 45
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|