35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from motrix_rl.rslrl.torch.train.dreamwaq_ppo import DreamWaQPPO
|
|
from motrix_rl.rslrl.torch.wrap_vec_env import RslrlNpEnvWrap
|
|
|
|
|
|
def test_dreamwaq_adaptive_learning_rate_is_capped_at_initial_rate():
|
|
assert DreamWaQPPO.max_learning_rate == 1e-3
|
|
|
|
|
|
def test_step_preserves_rewards_and_marks_only_truncations_as_timeouts():
|
|
state = SimpleNamespace(
|
|
obs=np.zeros((2, 3), dtype=np.float32),
|
|
reward=np.array([1.25, -0.5], dtype=np.float32),
|
|
done=np.array([True, True]),
|
|
terminated=np.array([True, False]),
|
|
truncated=np.array([False, True]),
|
|
info={},
|
|
)
|
|
wrapper = RslrlNpEnvWrap.__new__(RslrlNpEnvWrap)
|
|
wrapper._env = SimpleNamespace(step=lambda actions: state)
|
|
wrapper._device = torch.device("cpu")
|
|
wrapper._state = None
|
|
wrapper._num_envs = 2
|
|
wrapper.episode_length_buf = torch.zeros(2, dtype=torch.long)
|
|
|
|
_, rewards, dones, extras = wrapper.step(torch.zeros((2, 1)))
|
|
|
|
torch.testing.assert_close(rewards, torch.tensor([1.25, -0.5]))
|
|
torch.testing.assert_close(dones, torch.tensor([1.0, 1.0]))
|
|
torch.testing.assert_close(extras["time_outs"], torch.tensor([0.0, 1.0]))
|