File size: 3,049 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import pytest
import torch
import copy
from unittest.mock import patch
from ding.framework import OnlineRLContext, task
from ding.framework.middleware import TransitionList, inferencer, rolloutor
from ding.framework.middleware import StepCollector, EpisodeCollector
from ding.framework.middleware.tests import MockPolicy, MockEnv, CONFIG
@pytest.mark.unittest
def test_inferencer():
ctx = OnlineRLContext()
with patch("ding.policy.Policy", MockPolicy), patch("ding.envs.BaseEnvManagerV2", MockEnv):
policy = MockPolicy()
env = MockEnv()
inferencer(0, policy, env)(ctx)
assert isinstance(ctx.inference_output, dict)
assert ctx.inference_output[0] == {'action': torch.Tensor([0.])} # sum of zeros([2, 2])
assert ctx.inference_output[1] == {'action': torch.Tensor([4.])} # sum of ones([2, 2])
@pytest.mark.unittest
def test_rolloutor():
ctx = OnlineRLContext()
transitions = TransitionList(2)
with patch("ding.policy.Policy", MockPolicy), patch("ding.envs.BaseEnvManagerV2", MockEnv):
policy = MockPolicy()
env = MockEnv()
for _ in range(10):
inferencer(0, policy, env)(ctx)
rolloutor(policy, env, transitions)(ctx)
assert ctx.env_episode == 20 # 10 * env_num
assert ctx.env_step == 20 # 10 * env_num
@pytest.mark.unittest
def test_step_collector():
cfg = copy.deepcopy(CONFIG)
ctx = OnlineRLContext()
# test no random_collect_size
with patch("ding.policy.Policy", MockPolicy), patch("ding.envs.BaseEnvManagerV2", MockEnv):
with task.start():
policy = MockPolicy()
env = MockEnv()
collector = StepCollector(cfg, policy, env)
collector(ctx)
assert len(ctx.trajectories) == 16
assert ctx.trajectory_end_idx == [7, 15]
# test with random_collect_size
with patch("ding.policy.Policy", MockPolicy), patch("ding.envs.BaseEnvManagerV2", MockEnv):
with task.start():
policy = MockPolicy()
env = MockEnv()
collector = StepCollector(cfg, policy, env, random_collect_size=8)
collector(ctx)
assert len(ctx.trajectories) == 16
assert ctx.trajectory_end_idx == [7, 15]
@pytest.mark.unittest
def test_episode_collector():
cfg = copy.deepcopy(CONFIG)
ctx = OnlineRLContext()
# test no random_collect_size
with patch("ding.policy.Policy", MockPolicy), patch("ding.envs.BaseEnvManagerV2", MockEnv):
with task.start():
policy = MockPolicy()
env = MockEnv()
collector = EpisodeCollector(cfg, policy, env)
collector(ctx)
assert len(ctx.episodes) == 16
# test with random_collect_size
with patch("ding.policy.Policy", MockPolicy), patch("ding.envs.BaseEnvManagerV2", MockEnv):
with task.start():
policy = MockPolicy()
env = MockEnv()
collector = EpisodeCollector(cfg, policy, env, random_collect_size=8)
collector(ctx)
assert len(ctx.episodes) == 16
|