mirror of
https://github.com/vale981/ray
synced 2025-03-09 12:56:46 -04:00
60d4d5e1aa
* Remove all __future__ imports from RLlib. * Remove (object) again from tf_run_builder.py::TFRunBuilder. * Fix 2xLINT warnings. * Fix broken appo_policy import (must be appo_tf_policy) * Remove future imports from all other ray files (not just RLlib). * Remove future imports from all other ray files (not just RLlib). * Remove future import blocks that contain `unicode_literals` as well. Revert appo_tf_policy.py to appo_policy.py (belongs to another PR). * Add two empty lines before Schedule class. * Put back __future__ imports into determine_tests_to_run.py. Fails otherwise on a py2/print related error.
103 lines
2.6 KiB
Python
103 lines
2.6 KiB
Python
import numpy as np
|
|
|
|
from ray.rllib.optimizers.segment_tree import SumSegmentTree, MinSegmentTree
|
|
|
|
|
|
def test_tree_set():
|
|
tree = SumSegmentTree(4)
|
|
|
|
tree[2] = 1.0
|
|
tree[3] = 3.0
|
|
|
|
assert np.isclose(tree.sum(), 4.0)
|
|
assert np.isclose(tree.sum(0, 2), 0.0)
|
|
assert np.isclose(tree.sum(0, 3), 1.0)
|
|
assert np.isclose(tree.sum(2, 3), 1.0)
|
|
assert np.isclose(tree.sum(2, -1), 1.0)
|
|
assert np.isclose(tree.sum(2, 4), 4.0)
|
|
|
|
|
|
def test_tree_set_overlap():
|
|
tree = SumSegmentTree(4)
|
|
|
|
tree[2] = 1.0
|
|
tree[2] = 3.0
|
|
|
|
assert np.isclose(tree.sum(), 3.0)
|
|
assert np.isclose(tree.sum(2, 3), 3.0)
|
|
assert np.isclose(tree.sum(2, -1), 3.0)
|
|
assert np.isclose(tree.sum(2, 4), 3.0)
|
|
assert np.isclose(tree.sum(1, 2), 0.0)
|
|
|
|
|
|
def test_prefixsum_idx():
|
|
tree = SumSegmentTree(4)
|
|
|
|
tree[2] = 1.0
|
|
tree[3] = 3.0
|
|
|
|
assert tree.find_prefixsum_idx(0.0) == 2
|
|
assert tree.find_prefixsum_idx(0.5) == 2
|
|
assert tree.find_prefixsum_idx(0.99) == 2
|
|
assert tree.find_prefixsum_idx(1.01) == 3
|
|
assert tree.find_prefixsum_idx(3.00) == 3
|
|
assert tree.find_prefixsum_idx(4.00) == 3
|
|
|
|
|
|
def test_prefixsum_idx2():
|
|
tree = SumSegmentTree(4)
|
|
|
|
tree[0] = 0.5
|
|
tree[1] = 1.0
|
|
tree[2] = 1.0
|
|
tree[3] = 3.0
|
|
|
|
assert tree.find_prefixsum_idx(0.00) == 0
|
|
assert tree.find_prefixsum_idx(0.55) == 1
|
|
assert tree.find_prefixsum_idx(0.99) == 1
|
|
assert tree.find_prefixsum_idx(1.51) == 2
|
|
assert tree.find_prefixsum_idx(3.00) == 3
|
|
assert tree.find_prefixsum_idx(5.50) == 3
|
|
|
|
|
|
def test_max_interval_tree():
|
|
tree = MinSegmentTree(4)
|
|
|
|
tree[0] = 1.0
|
|
tree[2] = 0.5
|
|
tree[3] = 3.0
|
|
|
|
assert np.isclose(tree.min(), 0.5)
|
|
assert np.isclose(tree.min(0, 2), 1.0)
|
|
assert np.isclose(tree.min(0, 3), 0.5)
|
|
assert np.isclose(tree.min(0, -1), 0.5)
|
|
assert np.isclose(tree.min(2, 4), 0.5)
|
|
assert np.isclose(tree.min(3, 4), 3.0)
|
|
|
|
tree[2] = 0.7
|
|
|
|
assert np.isclose(tree.min(), 0.7)
|
|
assert np.isclose(tree.min(0, 2), 1.0)
|
|
assert np.isclose(tree.min(0, 3), 0.7)
|
|
assert np.isclose(tree.min(0, -1), 0.7)
|
|
assert np.isclose(tree.min(2, 4), 0.7)
|
|
assert np.isclose(tree.min(3, 4), 3.0)
|
|
|
|
tree[2] = 4.0
|
|
|
|
assert np.isclose(tree.min(), 1.0)
|
|
assert np.isclose(tree.min(0, 2), 1.0)
|
|
assert np.isclose(tree.min(0, 3), 1.0)
|
|
assert np.isclose(tree.min(0, -1), 1.0)
|
|
assert np.isclose(tree.min(2, 4), 3.0)
|
|
assert np.isclose(tree.min(2, 3), 4.0)
|
|
assert np.isclose(tree.min(2, -1), 4.0)
|
|
assert np.isclose(tree.min(3, 4), 3.0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_tree_set()
|
|
test_tree_set_overlap()
|
|
test_prefixsum_idx()
|
|
test_prefixsum_idx2()
|
|
test_max_interval_tree()
|