|
import pytest |
|
import torch |
|
from ding.framework import OnlineRLContext |
|
from ding.data.buffer import DequeBuffer |
|
from typing import Any |
|
import numpy as np |
|
import copy |
|
from ding.framework.middleware.functional.enhancer import reward_estimator, her_data_enhancer |
|
from unittest.mock import Mock, patch |
|
from ding.framework.middleware.tests import MockHerRewardModel, CONFIG |
|
|
|
DATA = [{'obs': torch.rand(2, 2), 'next_obs': torch.rand(2, 2)} for _ in range(20)] |
|
|
|
|
|
class MockRewardModel(Mock): |
|
|
|
def estimate(self, data: list) -> Any: |
|
assert len(data) == len(DATA) |
|
assert torch.equal(data[0]['obs'], DATA[0]['obs']) |
|
|
|
|
|
@pytest.mark.unittest |
|
def test_reward_estimator(): |
|
ctx = OnlineRLContext() |
|
ctx.train_data = copy.deepcopy(DATA) |
|
with patch("ding.reward_model.HerRewardModel", MockHerRewardModel): |
|
reward_estimator(cfg=None, reward_model=MockRewardModel())(ctx) |
|
|
|
|
|
@pytest.mark.unittest |
|
def test_her_data_enhancer(): |
|
cfg = copy.deepcopy(CONFIG) |
|
ctx = OnlineRLContext() |
|
|
|
with patch("ding.reward_model.HerRewardModel", MockHerRewardModel): |
|
mock_her_reward_model = MockHerRewardModel() |
|
buffer = DequeBuffer(mock_her_reward_model.episode_size) |
|
|
|
train_data = [ |
|
[ |
|
{ |
|
'action': torch.randint(low=0, high=5, size=(1, )), |
|
'collect_train_iter': torch.tensor([0]), |
|
'done': torch.tensor(False), |
|
'next_obs': torch.randint(low=0, high=2, size=(10, ), dtype=torch.float32), |
|
'obs': torch.randint(low=0, high=2, size=(10, ), dtype=torch.float32), |
|
'reward': torch.randint(low=0, high=2, size=(1, ), dtype=torch.float32), |
|
} for _ in range(np.random.choice([1, 4, 5], size=1)[0]) |
|
] for _ in range(mock_her_reward_model.episode_size) |
|
] |
|
|
|
for d in train_data: |
|
buffer.push(d) |
|
|
|
her_data_enhancer(cfg=cfg, buffer_=buffer, her_reward_model=MockHerRewardModel())(ctx) |
|
assert len(ctx.train_data) == mock_her_reward_model.episode_size * mock_her_reward_model.episode_element_size |
|
assert len(ctx.train_data[0]) == 6 |
|
|
|
buffer = DequeBuffer(cfg.policy.learn.batch_size) |
|
for d in train_data: |
|
buffer.push(d) |
|
mock_her_reward_model.episode_size = None |
|
her_data_enhancer(cfg=cfg, buffer_=buffer, her_reward_model=MockHerRewardModel())(ctx) |
|
assert len(ctx.train_data) == cfg.policy.learn.batch_size * mock_her_reward_model.episode_element_size |
|
assert len(ctx.train_data[0]) == 6 |
|
|
