diff --git a/tests/jax/test_jax_model_instantiators.py b/tests/jax/test_jax_model_instantiators.py index 51cd2aaf..96e044d4 100644 --- a/tests/jax/test_jax_model_instantiators.py +++ b/tests/jax/test_jax_model_instantiators.py @@ -11,23 +11,27 @@ from skrl.utils.model_instantiators.jax import Shape, categorical_model, deterministic_model, gaussian_model -@hypothesis.given(observation_space_size=st.integers(min_value=1, max_value=10), - action_space_size=st.integers(min_value=1, max_value=10)) +@hypothesis.given( + observation_space_size=st.integers(min_value=1, max_value=10), + action_space_size=st.integers(min_value=1, max_value=10), +) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_categorical_model(capsys, observation_space_size, action_space_size, device): observation_space = gym.spaces.Box(np.array([-1] * observation_space_size), np.array([1] * observation_space_size)) action_space = gym.spaces.Discrete(action_space_size) # TODO: randomize all parameters - model = categorical_model(observation_space=observation_space, - action_space=action_space, - device=device, - unnormalized_log_prob=True, - input_shape=Shape.STATES, - hiddens=[256, 256], - hidden_activation=["relu", "relu"], - output_shape=Shape.ACTIONS, - output_activation=None) + model = categorical_model( + observation_space=observation_space, + action_space=action_space, + device=device, + unnormalized_log_prob=True, + input_shape=Shape.STATES, + hiddens=[256, 256], + hidden_activation=["relu", "relu"], + output_shape=Shape.ACTIONS, + output_activation=None, + ) model.init_state_dict("model") with jax.default_device(model.device): @@ -35,24 +39,29 @@ def test_categorical_model(capsys, observation_space_size, action_space_size, de output = model.act({"states": observations}) assert output[0].shape == (10, 1) -@hypothesis.given(observation_space_size=st.integers(min_value=1, max_value=10), - action_space_size=st.integers(min_value=1, max_value=10)) + +@hypothesis.given( + observation_space_size=st.integers(min_value=1, max_value=10), + action_space_size=st.integers(min_value=1, max_value=10), +) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_deterministic_model(capsys, observation_space_size, action_space_size, device): observation_space = gym.spaces.Box(np.array([-1] * observation_space_size), np.array([1] * observation_space_size)) action_space = gym.spaces.Box(np.array([-1] * action_space_size), np.array([1] * action_space_size)) # TODO: randomize all parameters - model = deterministic_model(observation_space=observation_space, - action_space=action_space, - device=device, - clip_actions=False, - input_shape=Shape.STATES, - hiddens=[256, 256], - hidden_activation=["relu", "relu"], - output_shape=Shape.ACTIONS, - output_activation=None, - output_scale=1) + model = deterministic_model( + observation_space=observation_space, + action_space=action_space, + device=device, + clip_actions=False, + input_shape=Shape.STATES, + hiddens=[256, 256], + hidden_activation=["relu", "relu"], + output_shape=Shape.ACTIONS, + output_activation=None, + output_scale=1, + ) model.init_state_dict("model") with jax.default_device(model.device): @@ -60,28 +69,33 @@ def test_deterministic_model(capsys, observation_space_size, action_space_size, output = model.act({"states": observations}) assert output[0].shape == (10, model.num_actions) -@hypothesis.given(observation_space_size=st.integers(min_value=1, max_value=10), - action_space_size=st.integers(min_value=1, max_value=10)) + +@hypothesis.given( + observation_space_size=st.integers(min_value=1, max_value=10), + action_space_size=st.integers(min_value=1, max_value=10), +) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_gaussian_model(capsys, observation_space_size, action_space_size, device): observation_space = gym.spaces.Box(np.array([-1] * observation_space_size), np.array([1] * observation_space_size)) action_space = gym.spaces.Box(np.array([-1] * action_space_size), np.array([1] * action_space_size)) # TODO: randomize all parameters - model = gaussian_model(observation_space=observation_space, - action_space=action_space, - device=device, - clip_actions=False, - clip_log_std=True, - min_log_std=-20, - max_log_std=2, - initial_log_std=0, - input_shape=Shape.STATES, - hiddens=[256, 256], - hidden_activation=["relu", "relu"], - output_shape=Shape.ACTIONS, - output_activation=None, - output_scale=1) + model = gaussian_model( + observation_space=observation_space, + action_space=action_space, + device=device, + clip_actions=False, + clip_log_std=True, + min_log_std=-20, + max_log_std=2, + initial_log_std=0, + input_shape=Shape.STATES, + hiddens=[256, 256], + hidden_activation=["relu", "relu"], + output_shape=Shape.ACTIONS, + output_activation=None, + output_scale=1, + ) model.init_state_dict("model") with jax.default_device(model.device): diff --git a/tests/jax/test_jax_model_instantiators_definition.py b/tests/jax/test_jax_model_instantiators_definition.py index f6992394..6c2f3de9 100644 --- a/tests/jax/test_jax_model_instantiators_definition.py +++ b/tests/jax/test_jax_model_instantiators_definition.py @@ -22,6 +22,7 @@ def test_get_activation_function(capsys): assert activation is not None, f"{item} -> None" exec(f"{activation}(x)", _globals, {}) + def test_parse_input(capsys): # check for Shape enum (compatibility with prior versions) for input in [Shape.STATES, Shape.OBSERVATIONS, Shape.ACTIONS, Shape.STATES_ACTIONS]: @@ -43,6 +44,7 @@ def test_parse_input(capsys): output = _parse_input(str(input)) assert output.replace("'", '"') == statement, f"'{output}' != '{statement}'" + def test_generate_modules(capsys): _globals = {"nn": flax.linen} @@ -138,6 +140,7 @@ def test_generate_modules(capsys): assert isinstance(container, flax.linen.Sequential) assert len(container.layers) == 2 + def test_gaussian_model(capsys): device = "cpu" observation_space = gym.spaces.Box(np.array([-1] * 5), np.array([1] * 5)) @@ -161,19 +164,15 @@ def test_gaussian_model(capsys): """ content = yaml.safe_load(content) # source - model = gaussian_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=True, - **content) + model = gaussian_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=True, **content + ) with capsys.disabled(): print(model) # instance - model = gaussian_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=False, - **content) + model = gaussian_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=False, **content + ) model.init_state_dict("model") with capsys.disabled(): print(model) @@ -182,6 +181,7 @@ def test_gaussian_model(capsys): output = model.act({"states": observations}) assert output[0].shape == (10, 2) + def test_deterministic_model(capsys): device = "cpu" observation_space = gym.spaces.Box(np.array([-1] * 5), np.array([1] * 5)) @@ -202,19 +202,15 @@ def test_deterministic_model(capsys): """ content = yaml.safe_load(content) # source - model = deterministic_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=True, - **content) + model = deterministic_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=True, **content + ) with capsys.disabled(): print(model) # instance - model = deterministic_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=False, - **content) + model = deterministic_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=False, **content + ) model.init_state_dict("model") with capsys.disabled(): print(model) @@ -223,6 +219,7 @@ def test_deterministic_model(capsys): output = model.act({"states": observations}) assert output[0].shape == (10, 3) + def test_categorical_model(capsys): device = "cpu" observation_space = gym.spaces.Box(np.array([-1] * 5), np.array([1] * 5)) @@ -242,19 +239,15 @@ def test_categorical_model(capsys): """ content = yaml.safe_load(content) # source - model = categorical_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=True, - **content) + model = categorical_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=True, **content + ) with capsys.disabled(): print(model) # instance - model = categorical_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=False, - **content) + model = categorical_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=False, **content + ) model.init_state_dict("model") with capsys.disabled(): print(model) diff --git a/tests/jax/test_jax_utils_spaces.py b/tests/jax/test_jax_utils_spaces.py index 44d5d983..2a845193 100644 --- a/tests/jax/test_jax_utils_spaces.py +++ b/tests/jax/test_jax_utils_spaces.py @@ -15,7 +15,7 @@ sample_space, tensorize_space, unflatten_tensorized_space, - untensorize_space + untensorize_space, ) from ..stategies import gym_space_stategy, gymnasium_space_stategy @@ -29,6 +29,7 @@ def _check_backend(x, backend): else: raise ValueError(f"Invalid backend type: {backend}") + def check_sampled_space(space, x, n, backend): if isinstance(space, gymnasium.spaces.Box): _check_backend(x, backend) @@ -66,6 +67,7 @@ def occupied_size(s): space_size = compute_space_size(space, occupied_size=True) assert space_size == occupied_size(space) + @hypothesis.given(space=gymnasium_space_stategy()) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_tensorize_space(capsys, space: gymnasium.spaces.Space): @@ -97,6 +99,7 @@ def check_tensorized_space(s, x, n): tensorized_space = tensorize_space(space, sampled_space) check_tensorized_space(space, tensorized_space, 5) + @hypothesis.given(space=gymnasium_space_stategy()) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_untensorize_space(capsys, space: gymnasium.spaces.Space): @@ -108,7 +111,9 @@ def check_untensorized_space(s, x, squeeze_batch_dimension): assert isinstance(x, (np.ndarray, int)) assert isinstance(x, int) if squeeze_batch_dimension else x.shape == (1, 1) elif isinstance(s, gymnasium.spaces.MultiDiscrete): - assert isinstance(x, np.ndarray) and x.shape == s.nvec.shape if squeeze_batch_dimension else (1, *s.nvec.shape) + assert ( + isinstance(x, np.ndarray) and x.shape == s.nvec.shape if squeeze_batch_dimension else (1, *s.nvec.shape) + ) elif isinstance(s, gymnasium.spaces.Dict): list(map(check_untensorized_space, s.values(), x.values(), [squeeze_batch_dimension] * len(s))) elif isinstance(s, gymnasium.spaces.Tuple): @@ -124,6 +129,7 @@ def check_untensorized_space(s, x, squeeze_batch_dimension): untensorized_space = untensorize_space(space, tensorized_space, squeeze_batch_dimension=True) check_untensorized_space(space, untensorized_space, squeeze_batch_dimension=True) + @hypothesis.given(space=gymnasium_space_stategy(), batch_size=st.integers(min_value=1, max_value=10)) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_sample_space(capsys, space: gymnasium.spaces.Space, batch_size: int): @@ -134,6 +140,7 @@ def test_sample_space(capsys, space: gymnasium.spaces.Space, batch_size: int): sampled_space = sample_space(space, batch_size, backend="jax") check_sampled_space(space, sampled_space, batch_size, backend="jax") + @hypothesis.given(space=gymnasium_space_stategy()) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_flatten_tensorized_space(capsys, space: gymnasium.spaces.Space): @@ -147,6 +154,7 @@ def test_flatten_tensorized_space(capsys, space: gymnasium.spaces.Space): flattened_space = flatten_tensorized_space(tensorized_space) assert flattened_space.shape == (5, space_size) + @hypothesis.given(space=gymnasium_space_stategy()) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_unflatten_tensorized_space(capsys, space: gymnasium.spaces.Space): @@ -160,6 +168,7 @@ def test_unflatten_tensorized_space(capsys, space: gymnasium.spaces.Space): unflattened_space = unflatten_tensorized_space(space, flattened_space) check_sampled_space(space, unflattened_space, 5, backend="jax") + @hypothesis.given(space=gym_space_stategy()) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_convert_gym_space(capsys, space: gym.spaces.Space): diff --git a/tests/jax/test_jax_wrapper_gym.py b/tests/jax/test_jax_wrapper_gym.py index 7d7042ff..d2ef83cb 100644 --- a/tests/jax/test_jax_wrapper_gym.py +++ b/tests/jax/test_jax_wrapper_gym.py @@ -53,6 +53,7 @@ def test_env(capsys: pytest.CaptureFixture, backend: str): env.close() + @pytest.mark.parametrize("backend", ["jax", "numpy"]) @pytest.mark.parametrize("vectorization_mode", ["async", "sync"]) def test_vectorized_env(capsys: pytest.CaptureFixture, backend: str, vectorization_mode: str): diff --git a/tests/jax/test_jax_wrapper_gymnasium.py b/tests/jax/test_jax_wrapper_gymnasium.py index f8b47d5b..b7676d3f 100644 --- a/tests/jax/test_jax_wrapper_gymnasium.py +++ b/tests/jax/test_jax_wrapper_gymnasium.py @@ -52,6 +52,7 @@ def test_env(capsys: pytest.CaptureFixture, backend: str): env.close() + @pytest.mark.parametrize("backend", ["jax", "numpy"]) @pytest.mark.parametrize("vectorization_mode", ["async", "sync"]) def test_vectorized_env(capsys: pytest.CaptureFixture, backend: str, vectorization_mode: str): diff --git a/tests/jax/test_jax_wrapper_isaacgym.py b/tests/jax/test_jax_wrapper_isaacgym.py index 4629e666..d129bc02 100644 --- a/tests/jax/test_jax_wrapper_isaacgym.py +++ b/tests/jax/test_jax_wrapper_isaacgym.py @@ -30,7 +30,7 @@ def __init__(self, num_states) -> None: self.state_space = gym.spaces.Box(np.ones(self.num_states) * -np.Inf, np.ones(self.num_states) * np.Inf) self.observation_space = gym.spaces.Box(np.ones(self.num_obs) * -np.Inf, np.ones(self.num_obs) * np.Inf) - self.action_space = gym.spaces.Box(np.ones(self.num_actions) * -1., np.ones(self.num_actions) * 1.) + self.action_space = gym.spaces.Box(np.ones(self.num_actions) * -1.0, np.ones(self.num_actions) * 1.0) def reset(self) -> Dict[str, torch.Tensor]: obs_dict = {} diff --git a/tests/jax/test_jax_wrapper_isaaclab.py b/tests/jax/test_jax_wrapper_isaaclab.py index 9b62c5e4..b454fd2d 100644 --- a/tests/jax/test_jax_wrapper_isaaclab.py +++ b/tests/jax/test_jax_wrapper_isaaclab.py @@ -26,6 +26,7 @@ Dict[AgentID, dict], ] + class IsaacLabEnv(gym.Env): def __init__(self, num_states) -> None: self.num_actions = 1 @@ -61,7 +62,9 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) def step(self, action: torch.Tensor) -> VecEnvStepReturn: assert action.clone().shape == torch.Size([self.num_envs, 1]) - observations = {"policy": torch.ones((self.num_envs, self.num_observations), device=self.device, dtype=torch.float32)} + observations = { + "policy": torch.ones((self.num_envs, self.num_observations), device=self.device, dtype=torch.float32) + } rewards = torch.zeros(self.num_envs, device=self.device, dtype=torch.float32) terminated = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) truncated = torch.zeros_like(terminated) @@ -102,9 +105,7 @@ def _configure_env_spaces(self): if not self.num_states: self.state_space = None if self.num_states < 0: - self.state_space = gym.spaces.Box( - low=-np.inf, high=np.inf, shape=(sum(self.num_observations.values()),) - ) + self.state_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(sum(self.num_observations.values()),)) else: self.state_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.num_states,)) @@ -112,16 +113,28 @@ def _configure_env_spaces(self): def unwrapped(self): return self - def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Dict[AgentID, ObsType], dict]: - observations = {agent: torch.ones((self.num_envs, self.num_observations[agent]), device=self.device) for agent in self.possible_agents} + def reset( + self, seed: int | None = None, options: dict[str, Any] | None = None + ) -> tuple[Dict[AgentID, ObsType], dict]: + observations = { + agent: torch.ones((self.num_envs, self.num_observations[agent]), device=self.device) + for agent in self.possible_agents + } return observations, self.extras def step(self, action: Dict[AgentID, torch.Tensor]) -> EnvStepReturn: for agent in self.possible_agents: assert action[agent].clone().shape == torch.Size([self.num_envs, self.num_actions[agent]]) - observations = {agent: torch.ones((self.num_envs, self.num_observations[agent]), device=self.device) for agent in self.possible_agents} - rewards = {agent: torch.zeros(self.num_envs, device=self.device, dtype=torch.float32) for agent in self.possible_agents} - terminated = {agent: torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) for agent in self.possible_agents} + observations = { + agent: torch.ones((self.num_envs, self.num_observations[agent]), device=self.device) + for agent in self.possible_agents + } + rewards = { + agent: torch.zeros(self.num_envs, device=self.device, dtype=torch.float32) for agent in self.possible_agents + } + terminated = { + agent: torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) for agent in self.possible_agents + } truncated = {agent: torch.zeros_like(terminated[agent]) for agent in self.possible_agents} return observations, rewards, terminated, truncated, self.extras @@ -183,6 +196,7 @@ def test_env(capsys: pytest.CaptureFixture, backend: str, num_states: int): env.close() + @pytest.mark.parametrize("backend", ["jax", "numpy"]) @pytest.mark.parametrize("num_states", [0, 5]) def test_multi_agent_env(capsys: pytest.CaptureFixture, backend: str, num_states: int): @@ -192,7 +206,10 @@ def test_multi_agent_env(capsys: pytest.CaptureFixture, backend: str, num_states num_envs = 10 num_agents = 3 possible_agents = [f"agent_{i}" for i in range(num_agents)] - action = {f"agent_{i}": jnp.ones((num_envs, i + 10)) if backend == "jax" else np.ones((num_envs, i + 10)) for i in range(num_agents)} + action = { + f"agent_{i}": jnp.ones((num_envs, i + 10)) if backend == "jax" else np.ones((num_envs, i + 10)) + for i in range(num_agents) + } # load wrap the environment original_env = IsaacLabMultiAgentEnv(num_states) diff --git a/tests/jax/test_jax_wrapper_omniisaacgym.py b/tests/jax/test_jax_wrapper_omniisaacgym.py index 5f59014d..caf8e606 100644 --- a/tests/jax/test_jax_wrapper_omniisaacgym.py +++ b/tests/jax/test_jax_wrapper_omniisaacgym.py @@ -28,9 +28,16 @@ def __init__(self, num_states) -> None: self.device = "cpu" # initialize data spaces (defaults to gym.Box) - self.action_space = gym.spaces.Box(np.ones(self.num_actions, dtype=np.float32) * -1.0, np.ones(self.num_actions, dtype=np.float32) * 1.0) - self.observation_space = gym.spaces.Box(np.ones(self.num_observations, dtype=np.float32) * -np.Inf, np.ones(self.num_observations, dtype=np.float32) * np.Inf) - self.state_space = gym.spaces.Box(np.ones(self.num_states, dtype=np.float32) * -np.Inf, np.ones(self.num_states, dtype=np.float32) * np.Inf) + self.action_space = gym.spaces.Box( + np.ones(self.num_actions, dtype=np.float32) * -1.0, np.ones(self.num_actions, dtype=np.float32) * 1.0 + ) + self.observation_space = gym.spaces.Box( + np.ones(self.num_observations, dtype=np.float32) * -np.Inf, + np.ones(self.num_observations, dtype=np.float32) * np.Inf, + ) + self.state_space = gym.spaces.Box( + np.ones(self.num_states, dtype=np.float32) * -np.Inf, np.ones(self.num_states, dtype=np.float32) * np.Inf + ) def reset(self): observations = {"obs": torch.ones((self.num_envs, self.num_observations), device=self.device)} @@ -38,7 +45,9 @@ def reset(self): def step(self, actions): assert actions.clone().shape == torch.Size([self.num_envs, 1]) - observations = {"obs": torch.ones((self.num_envs, self.num_observations), device=self.device, dtype=torch.float32)} + observations = { + "obs": torch.ones((self.num_envs, self.num_observations), device=self.device, dtype=torch.float32) + } rewards = torch.zeros(self.num_envs, device=self.device, dtype=torch.float32) dones = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) return observations, rewards, dones, self.extras diff --git a/tests/jax/test_jax_wrapper_pettingzoo.py b/tests/jax/test_jax_wrapper_pettingzoo.py index d62756fe..93c77df3 100644 --- a/tests/jax/test_jax_wrapper_pettingzoo.py +++ b/tests/jax/test_jax_wrapper_pettingzoo.py @@ -21,7 +21,10 @@ def test_env(capsys: pytest.CaptureFixture, backend: str): num_envs = 1 num_agents = 20 possible_agents = [f"piston_{i}" for i in range(num_agents)] - action = {f"piston_{i}": jnp.ones((num_envs, 1)) if backend == "jax" else np.ones((num_envs, 1)) for i in range(num_agents)} + action = { + f"piston_{i}": jnp.ones((num_envs, 1)) if backend == "jax" else np.ones((num_envs, 1)) + for i in range(num_agents) + } # load wrap the environment original_env = pistonball_v6.parallel_env(n_pistons=num_agents, continuous=True, max_cycles=125) @@ -38,7 +41,11 @@ def test_env(capsys: pytest.CaptureFixture, backend: str): assert isinstance(env.action_spaces, Mapping) and len(env.action_spaces) == num_agents for agent in possible_agents: assert isinstance(env.state_space(agent), gym.Space) and env.state_space(agent).shape == (560, 880, 3) - assert isinstance(env.observation_space(agent), gym.Space) and env.observation_space(agent).shape == (457, 120, 3) + assert isinstance(env.observation_space(agent), gym.Space) and env.observation_space(agent).shape == ( + 457, + 120, + 3, + ) assert isinstance(env.action_space(agent), gym.Space) and env.action_space(agent).shape == (1,) assert isinstance(env.possible_agents, list) and sorted(env.possible_agents) == sorted(possible_agents) assert isinstance(env.num_envs, int) and env.num_envs == num_envs @@ -54,7 +61,10 @@ def test_env(capsys: pytest.CaptureFixture, backend: str): assert isinstance(observation, Mapping) assert isinstance(info, Mapping) for agent in possible_agents: - assert isinstance(observation[agent], Array) and observation[agent].shape == (num_envs, math.prod((457, 120, 3))) + assert isinstance(observation[agent], Array) and observation[agent].shape == ( + num_envs, + math.prod((457, 120, 3)), + ) for _ in range(3): observation, reward, terminated, truncated, info = env.step(action) state = env.state() @@ -65,7 +75,10 @@ def test_env(capsys: pytest.CaptureFixture, backend: str): assert isinstance(truncated, Mapping) assert isinstance(info, Mapping) for agent in possible_agents: - assert isinstance(observation[agent], Array) and observation[agent].shape == (num_envs, math.prod((457, 120, 3))) + assert isinstance(observation[agent], Array) and observation[agent].shape == ( + num_envs, + math.prod((457, 120, 3)), + ) assert isinstance(reward[agent], Array) and reward[agent].shape == (num_envs, 1) assert isinstance(terminated[agent], Array) and terminated[agent].shape == (num_envs, 1) assert isinstance(truncated[agent], Array) and truncated[agent].shape == (num_envs, 1) diff --git a/tests/stategies.py b/tests/stategies.py index 409b9771..4840a541 100644 --- a/tests/stategies.py +++ b/tests/stategies.py @@ -28,11 +28,14 @@ def gymnasium_space_stategy(draw, space_type: str = "", remaining_iterations: in return gymnasium.spaces.Dict(spaces) elif space_type == "Tuple": remaining_iterations -= 1 - spaces = draw(st.lists(gymnasium_space_stategy(remaining_iterations=remaining_iterations), min_size=1, max_size=3)) + spaces = draw( + st.lists(gymnasium_space_stategy(remaining_iterations=remaining_iterations), min_size=1, max_size=3) + ) return gymnasium.spaces.Tuple(spaces) else: raise ValueError(f"Invalid space type: {space_type}") + @st.composite def gym_space_stategy(draw, space_type: str = "", remaining_iterations: int = 5) -> gym.spaces.Space: if not space_type: diff --git a/tests/test_agents.py b/tests/test_agents.py index 2fcb36d5..48a6d3dd 100644 --- a/tests/test_agents.py +++ b/tests/test_agents.py @@ -23,36 +23,39 @@ @pytest.fixture def classes_and_kwargs(): - return [(A2C, {"models": {"policy": DummyModel()}}), - (AMP, {"models": {"policy": DummyModel()}}), - (CEM, {"models": {"policy": DummyModel()}}), - (DDPG, {"models": {"policy": DummyModel()}}), - (DQN, {"models": {"policy": DummyModel()}}), - (DDQN, {"models": {"policy": DummyModel()}}), - (PPO, {"models": {"policy": DummyModel()}}), - (Q_LEARNING, {"models": {"policy": DummyModel()}}), - (SAC, {"models": {"policy": DummyModel()}}), - (SARSA, {"models": {"policy": DummyModel()}}), - (TD3, {"models": {"policy": DummyModel()}}), - (TRPO, {"models": {"policy": DummyModel()}})] + return [ + (A2C, {"models": {"policy": DummyModel()}}), + (AMP, {"models": {"policy": DummyModel()}}), + (CEM, {"models": {"policy": DummyModel()}}), + (DDPG, {"models": {"policy": DummyModel()}}), + (DQN, {"models": {"policy": DummyModel()}}), + (DDQN, {"models": {"policy": DummyModel()}}), + (PPO, {"models": {"policy": DummyModel()}}), + (Q_LEARNING, {"models": {"policy": DummyModel()}}), + (SAC, {"models": {"policy": DummyModel()}}), + (SARSA, {"models": {"policy": DummyModel()}}), + (TD3, {"models": {"policy": DummyModel()}}), + (TRPO, {"models": {"policy": DummyModel()}}), + ] def test_agent(capsys, classes_and_kwargs): for klass, kwargs in classes_and_kwargs: - cfg = {"learning_starts": 1, - "experiment": {"write_interval": 0}} + cfg = {"learning_starts": 1, "experiment": {"write_interval": 0}} agent: Agent = klass(cfg=cfg, **kwargs) agent.init() agent.pre_interaction(timestep=0, timesteps=1) # agent.act(None, timestep=0, timestesps=1) - agent.record_transition(states=torch.tensor([]), - actions=torch.tensor([]), - rewards=torch.tensor([]), - next_states=torch.tensor([]), - terminated=torch.tensor([]), - truncated=torch.tensor([]), - infos={}, - timestep=0, - timesteps=1) + agent.record_transition( + states=torch.tensor([]), + actions=torch.tensor([]), + rewards=torch.tensor([]), + next_states=torch.tensor([]), + terminated=torch.tensor([]), + truncated=torch.tensor([]), + infos={}, + timestep=0, + timesteps=1, + ) agent.post_interaction(timestep=0, timesteps=1) diff --git a/tests/test_envs.py b/tests/test_envs.py index deac4c4d..00f5480c 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -15,8 +15,19 @@ def classes_and_kwargs(): return [] -@pytest.mark.parametrize("wrapper", ["gym", "gymnasium", "dm", "robosuite", \ - "isaacgym-preview2", "isaacgym-preview3", "isaacgym-preview4", "omniverse-isaacgym"]) +@pytest.mark.parametrize( + "wrapper", + [ + "gym", + "gymnasium", + "dm", + "robosuite", + "isaacgym-preview2", + "isaacgym-preview3", + "isaacgym-preview4", + "omniverse-isaacgym", + ], +) def test_wrap_env(capsys, classes_and_kwargs, wrapper): env = DummyEnv(num_envs=1) diff --git a/tests/test_examples_deepmind.py b/tests/test_examples_deepmind.py index 98222a3b..6cdf2a92 100644 --- a/tests/test_examples_deepmind.py +++ b/tests/test_examples_deepmind.py @@ -8,9 +8,10 @@ EXAMPLE_DIR = "deepmind" -SCRIPTS = ["dm_suite_cartpole_swingup_ddpg.py", - "dm_manipulation_stack_sac.py", ""] -EXAMPLES_DIR = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples")) +SCRIPTS = ["dm_suite_cartpole_swingup_ddpg.py", "dm_manipulation_stack_sac.py", ""] +EXAMPLES_DIR = os.path.abspath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples") +) COMMANDS = [f"python {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)}" for script in SCRIPTS] diff --git a/tests/test_examples_gym.py b/tests/test_examples_gym.py index 479fe841..73ff3a19 100644 --- a/tests/test_examples_gym.py +++ b/tests/test_examples_gym.py @@ -8,11 +8,10 @@ EXAMPLE_DIR = "gym" -SCRIPTS = ["ddpg_gym_pendulum.py", - "cem_gym_cartpole.py", - "dqn_gym_cartpole.py", - "q_learning_gym_frozen_lake.py"] -EXAMPLES_DIR = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples")) +SCRIPTS = ["ddpg_gym_pendulum.py", "cem_gym_cartpole.py", "dqn_gym_cartpole.py", "q_learning_gym_frozen_lake.py"] +EXAMPLES_DIR = os.path.abspath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples") +) COMMANDS = [f"python {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)}" for script in SCRIPTS] diff --git a/tests/test_examples_gymnasium.py b/tests/test_examples_gymnasium.py index a643a5a6..96b0a5a3 100644 --- a/tests/test_examples_gymnasium.py +++ b/tests/test_examples_gymnasium.py @@ -8,11 +8,15 @@ EXAMPLE_DIR = "gymnasium" -SCRIPTS = ["ddpg_gymnasium_pendulum.py", - "cem_gymnasium_cartpole.py", - "dqn_gymnasium_cartpole.py", - "q_learning_gymnasium_frozen_lake.py"] -EXAMPLES_DIR = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples")) +SCRIPTS = [ + "ddpg_gymnasium_pendulum.py", + "cem_gymnasium_cartpole.py", + "dqn_gymnasium_cartpole.py", + "q_learning_gymnasium_frozen_lake.py", +] +EXAMPLES_DIR = os.path.abspath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples") +) COMMANDS = [f"python {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)}" for script in SCRIPTS] diff --git a/tests/test_examples_isaac_orbit.py b/tests/test_examples_isaac_orbit.py index a9f17bee..8e89dfea 100644 --- a/tests/test_examples_isaac_orbit.py +++ b/tests/test_examples_isaac_orbit.py @@ -13,8 +13,13 @@ EXAMPLE_DIR = "isaacorbit" SCRIPTS = ["ppo_cartpole.py"] -EXAMPLES_DIR = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples")) -COMMANDS = [f"{PYTHON_ENVIRONMENT} {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)} --headless --num_envs 64" for script in SCRIPTS] +EXAMPLES_DIR = os.path.abspath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples") +) +COMMANDS = [ + f"{PYTHON_ENVIRONMENT} {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)} --headless --num_envs 64" + for script in SCRIPTS +] @pytest.mark.parametrize("command", COMMANDS) diff --git a/tests/test_examples_isaacgym.py b/tests/test_examples_isaacgym.py index 523a86db..2ca0f1e1 100644 --- a/tests/test_examples_isaacgym.py +++ b/tests/test_examples_isaacgym.py @@ -8,9 +8,10 @@ EXAMPLE_DIR = "isaacgym" -SCRIPTS = ["ppo_cartpole.py", - "trpo_cartpole.py"] -EXAMPLES_DIR = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples")) +SCRIPTS = ["ppo_cartpole.py", "trpo_cartpole.py"] +EXAMPLES_DIR = os.path.abspath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples") +) COMMANDS = [f"python {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)} headless=True num_envs=64" for script in SCRIPTS] diff --git a/tests/test_examples_isaacsim.py b/tests/test_examples_isaacsim.py index ef367110..5b1fdd26 100644 --- a/tests/test_examples_isaacsim.py +++ b/tests/test_examples_isaacsim.py @@ -13,7 +13,9 @@ EXAMPLE_DIR = "isaacsim" SCRIPTS = ["cartpole_example_skrl.py"] -EXAMPLES_DIR = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples")) +EXAMPLES_DIR = os.path.abspath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples") +) COMMANDS = [f"{PYTHON_ENVIRONMENT} {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)}" for script in SCRIPTS] diff --git a/tests/test_examples_omniisaacgym.py b/tests/test_examples_omniisaacgym.py index d92f7c3c..d0ccddf5 100644 --- a/tests/test_examples_omniisaacgym.py +++ b/tests/test_examples_omniisaacgym.py @@ -13,8 +13,13 @@ EXAMPLE_DIR = "omniisaacgym" SCRIPTS = ["ppo_cartpole.py"] -EXAMPLES_DIR = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples")) -COMMANDS = [f"{PYTHON_ENVIRONMENT} {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)} headless=True num_envs=64" for script in SCRIPTS] +EXAMPLES_DIR = os.path.abspath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples") +) +COMMANDS = [ + f"{PYTHON_ENVIRONMENT} {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)} headless=True num_envs=64" + for script in SCRIPTS +] @pytest.mark.parametrize("command", COMMANDS) diff --git a/tests/test_examples_robosuite.py b/tests/test_examples_robosuite.py index a03b2f61..c238538d 100644 --- a/tests/test_examples_robosuite.py +++ b/tests/test_examples_robosuite.py @@ -9,7 +9,9 @@ EXAMPLE_DIR = "robosuite" SCRIPTS = [] -EXAMPLES_DIR = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples")) +EXAMPLES_DIR = os.path.abspath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples") +) COMMANDS = [f"python {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)}" for script in SCRIPTS] diff --git a/tests/test_examples_shimmy.py b/tests/test_examples_shimmy.py index 3283d239..623ad53d 100644 --- a/tests/test_examples_shimmy.py +++ b/tests/test_examples_shimmy.py @@ -8,10 +8,14 @@ EXAMPLE_DIR = "shimmy" -SCRIPTS = ["dqn_shimmy_atari_pong.py", - "sac_shimmy_dm_control_acrobot_swingup_sparse.py", - "ddpg_openai_gym_compatibility_pendulum.py"] -EXAMPLES_DIR = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples")) +SCRIPTS = [ + "dqn_shimmy_atari_pong.py", + "sac_shimmy_dm_control_acrobot_swingup_sparse.py", + "ddpg_openai_gym_compatibility_pendulum.py", +] +EXAMPLES_DIR = os.path.abspath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "docs", "source", "examples") +) COMMANDS = [f"python {os.path.join(EXAMPLES_DIR, EXAMPLE_DIR, script)}" for script in SCRIPTS] diff --git a/tests/test_jax_memories_memory.py b/tests/test_jax_memories_memory.py index 586205b6..1f0d996f 100644 --- a/tests/test_jax_memories_memory.py +++ b/tests/test_jax_memories_memory.py @@ -46,13 +46,25 @@ def test_tensor_names(self): # test memory.get_tensor_by_name for name, size, dtype in zip(self.names, self.sizes, self.dtypes): tensor = memory.get_tensor_by_name(name, keepdim=True) - self.assertSequenceEqual(memory.get_tensor_by_name(name, keepdim=True).shape, (memory_size, num_envs, size), "get_tensor_by_name(..., keepdim=True)") - self.assertSequenceEqual(memory.get_tensor_by_name(name, keepdim=False).shape, (memory_size * num_envs, size), "get_tensor_by_name(..., keepdim=False)") - self.assertEqual(memory.get_tensor_by_name(name, keepdim=True).dtype, dtype, "get_tensor_by_name(...).dtype") + self.assertSequenceEqual( + memory.get_tensor_by_name(name, keepdim=True).shape, + (memory_size, num_envs, size), + "get_tensor_by_name(..., keepdim=True)", + ) + self.assertSequenceEqual( + memory.get_tensor_by_name(name, keepdim=False).shape, + (memory_size * num_envs, size), + "get_tensor_by_name(..., keepdim=False)", + ) + self.assertEqual( + memory.get_tensor_by_name(name, keepdim=True).dtype, dtype, "get_tensor_by_name(...).dtype" + ) # test memory.set_tensor_by_name for name, size, dtype in zip(self.names, self.sizes, self.raw_dtypes): - new_tensor = jnp.arange(memory_size * num_envs * size).reshape(memory_size, num_envs, size).astype(dtype) + new_tensor = ( + jnp.arange(memory_size * num_envs * size).reshape(memory_size, num_envs, size).astype(dtype) + ) memory.set_tensor_by_name(name, new_tensor) tensor = memory.get_tensor_by_name(name, keepdim=True) self.assertTrue((tensor == new_tensor).all().item(), "set_tensor_by_name(...)") @@ -68,30 +80,39 @@ def test_sample(self): # fill memory for name, size, dtype in zip(self.names, self.sizes, self.raw_dtypes): - new_tensor = jnp.arange(memory_size * num_envs * size).reshape(memory_size, num_envs, size).astype(dtype) + new_tensor = ( + jnp.arange(memory_size * num_envs * size).reshape(memory_size, num_envs, size).astype(dtype) + ) memory.set_tensor_by_name(name, new_tensor) # test memory.sample_all for i, mini_batches in enumerate(self.mini_batches): samples = memory.sample_all(self.names, mini_batches=mini_batches) for sample, name, size in zip(samples[i], self.names, self.sizes): - self.assertSequenceEqual(sample.shape, (memory_size * num_envs, size), f"sample_all(...).shape with mini_batches={mini_batches}") + self.assertSequenceEqual( + sample.shape, + (memory_size * num_envs, size), + f"sample_all(...).shape with mini_batches={mini_batches}", + ) tensor = memory.get_tensor_by_name(name, keepdim=True) - self.assertTrue((sample.reshape(memory_size, num_envs, size) == tensor).all().item(), f"sample_all(...) with mini_batches={mini_batches}") + self.assertTrue( + (sample.reshape(memory_size, num_envs, size) == tensor).all().item(), + f"sample_all(...) with mini_batches={mini_batches}", + ) -if __name__ == '__main__': +if __name__ == "__main__": import sys - if not sys.argv[-1] == '--debug': - raise RuntimeError('Test can only be runned manually with --debug flag') + if not sys.argv[-1] == "--debug": + raise RuntimeError("Test can only be runned manually with --debug flag") test = TestCase() test.setUp() for method in dir(test): - if method.startswith('test_'): - print('Running test: {}'.format(method)) + if method.startswith("test_"): + print("Running test: {}".format(method)) getattr(test, method)() test.tearDown() - print('All tests passed.') + print("All tests passed.") diff --git a/tests/test_memories.py b/tests/test_memories.py index 1cf2da2c..746daea5 100644 --- a/tests/test_memories.py +++ b/tests/test_memories.py @@ -17,7 +17,9 @@ def classes_and_kwargs(): @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_device(capsys, classes_and_kwargs, device): - _device = torch.device(device) if device is not None else torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + _device = ( + torch.device(device) if device is not None else torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + ) for klass, kwargs in classes_and_kwargs: try: @@ -30,7 +32,12 @@ def test_device(capsys, classes_and_kwargs, device): assert memory.device == _device # defined device -@hypothesis.given(names=st.sets(st.text(alphabet=string.ascii_letters + string.digits + "_", min_size=1, max_size=10), min_size=1, max_size=10)) + +@hypothesis.given( + names=st.sets( + st.text(alphabet=string.ascii_letters + string.digits + "_", min_size=1, max_size=10), min_size=1, max_size=10 + ) +) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_create_tensors(capsys, classes_and_kwargs, names): for klass, kwargs in classes_and_kwargs: @@ -41,9 +48,12 @@ def test_create_tensors(capsys, classes_and_kwargs, names): assert memory.get_tensor_names() == sorted(names) -@hypothesis.given(memory_size=st.integers(min_value=1, max_value=100), - num_envs=st.integers(min_value=1, max_value=10), - num_samples=st.integers(min_value=1, max_value=500)) + +@hypothesis.given( + memory_size=st.integers(min_value=1, max_value=100), + num_envs=st.integers(min_value=1, max_value=10), + num_samples=st.integers(min_value=1, max_value=500), +) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_add_samples(capsys, classes_and_kwargs, memory_size, num_envs, num_samples): for klass, kwargs in classes_and_kwargs: diff --git a/tests/test_model_instantiators.py b/tests/test_model_instantiators.py index 75ec2bf9..ab8756c7 100644 --- a/tests/test_model_instantiators.py +++ b/tests/test_model_instantiators.py @@ -11,16 +11,13 @@ categorical_model, deterministic_model, gaussian_model, - multivariate_gaussian_model + multivariate_gaussian_model, ) @pytest.fixture def classes_and_kwargs(): - return [(categorical_model, {}), - (deterministic_model, {}), - (gaussian_model, {}), - (multivariate_gaussian_model, {})] + return [(categorical_model, {}), (deterministic_model, {}), (gaussian_model, {}), (multivariate_gaussian_model, {})] def test_models(capsys, classes_and_kwargs): diff --git a/tests/test_resources_noises.py b/tests/test_resources_noises.py index 108da62e..927eb8b8 100644 --- a/tests/test_resources_noises.py +++ b/tests/test_resources_noises.py @@ -10,13 +10,17 @@ @pytest.fixture def classes_and_kwargs(): - return [(GaussianNoise, {"mean": 0, "std": 1}), - (OrnsteinUhlenbeckNoise, {"theta": 0.1, "sigma": 0.2, "base_scale": 0.3})] + return [ + (GaussianNoise, {"mean": 0, "std": 1}), + (OrnsteinUhlenbeckNoise, {"theta": 0.1, "sigma": 0.2, "base_scale": 0.3}), + ] @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_device(capsys, classes_and_kwargs, device): - _device = torch.device(device) if device is not None else torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + _device = ( + torch.device(device) if device is not None else torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + ) for klass, kwargs in classes_and_kwargs: try: @@ -31,6 +35,7 @@ def test_device(capsys, classes_and_kwargs, device): assert noise.device == _device # defined device assert output.device == _device # runtime device + @hypothesis.given(size=st.lists(st.integers(min_value=1, max_value=10), max_size=5)) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_sample(capsys, classes_and_kwargs, size): diff --git a/tests/test_resources_preprocessors.py b/tests/test_resources_preprocessors.py index 66f817cb..f340c198 100644 --- a/tests/test_resources_preprocessors.py +++ b/tests/test_resources_preprocessors.py @@ -18,7 +18,9 @@ def classes_and_kwargs(): @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_device(capsys, classes_and_kwargs, device): - _device = torch.device(device) if device is not None else torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + _device = ( + torch.device(device) if device is not None else torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + ) for klass, kwargs in classes_and_kwargs: try: @@ -32,10 +34,16 @@ def test_device(capsys, classes_and_kwargs, device): assert preprocessor.device == _device # defined device assert preprocessor(torch.ones(kwargs["size"], device=_device)).device == _device # runtime device -@pytest.mark.parametrize("space_and_size", [(gym.spaces.Box(low=-1, high=1, shape=(2, 3)), 6), - (gymnasium.spaces.Box(low=-1, high=1, shape=(2, 3)), 6), - (gym.spaces.Discrete(n=3), 1), - (gymnasium.spaces.Discrete(n=3), 1)]) + +@pytest.mark.parametrize( + "space_and_size", + [ + (gym.spaces.Box(low=-1, high=1, shape=(2, 3)), 6), + (gymnasium.spaces.Box(low=-1, high=1, shape=(2, 3)), 6), + (gym.spaces.Discrete(n=3), 1), + (gymnasium.spaces.Discrete(n=3), 1), + ], +) def test_forward(capsys, classes_and_kwargs, space_and_size): for klass, kwargs in classes_and_kwargs: space, size = space_and_size diff --git a/tests/test_resources_schedulers.py b/tests/test_resources_schedulers.py index 79b1db32..692710e4 100644 --- a/tests/test_resources_schedulers.py +++ b/tests/test_resources_schedulers.py @@ -13,8 +13,9 @@ def classes_and_kwargs(): return [(KLAdaptiveRL, {})] -@pytest.mark.parametrize("optimizer", [torch.optim.Adam([torch.ones((1,))], lr=0.1), - torch.optim.SGD([torch.ones((1,))], lr=0.1)]) +@pytest.mark.parametrize( + "optimizer", [torch.optim.Adam([torch.ones((1,))], lr=0.1), torch.optim.SGD([torch.ones((1,))], lr=0.1)] +) def test_step(capsys, classes_and_kwargs, optimizer): for klass, kwargs in classes_and_kwargs: scheduler = klass(optimizer, **kwargs) diff --git a/tests/test_trainers.py b/tests/test_trainers.py index f89c5a8b..f4e98265 100644 --- a/tests/test_trainers.py +++ b/tests/test_trainers.py @@ -12,9 +12,11 @@ @pytest.fixture def classes_and_kwargs(): - return [(ManualTrainer, {"cfg": {"timesteps": 100}}), - (ParallelTrainer, {"cfg": {"timesteps": 100}}), - (SequentialTrainer, {"cfg": {"timesteps": 100}})] + return [ + (ManualTrainer, {"cfg": {"timesteps": 100}}), + (ParallelTrainer, {"cfg": {"timesteps": 100}}), + (SequentialTrainer, {"cfg": {"timesteps": 100}}), + ] def test_train(capsys, classes_and_kwargs): @@ -26,6 +28,7 @@ def test_train(capsys, classes_and_kwargs): trainer.train() + def test_eval(capsys, classes_and_kwargs): env = DummyEnv(num_envs=1) agent = DummyAgent() diff --git a/tests/torch/test_torch_model_instantiators.py b/tests/torch/test_torch_model_instantiators.py index 44ff6c22..e84e5ec1 100644 --- a/tests/torch/test_torch_model_instantiators.py +++ b/tests/torch/test_torch_model_instantiators.py @@ -13,151 +13,175 @@ deterministic_model, gaussian_model, multivariate_gaussian_model, - shared_model + shared_model, ) -@hypothesis.given(observation_space_size=st.integers(min_value=1, max_value=10), - action_space_size=st.integers(min_value=1, max_value=10)) +@hypothesis.given( + observation_space_size=st.integers(min_value=1, max_value=10), + action_space_size=st.integers(min_value=1, max_value=10), +) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_categorical_model(capsys, observation_space_size, action_space_size, device): observation_space = gym.spaces.Box(np.array([-1] * observation_space_size), np.array([1] * observation_space_size)) action_space = gym.spaces.Discrete(action_space_size) # TODO: randomize all parameters - model = categorical_model(observation_space=observation_space, - action_space=action_space, - device=device, - unnormalized_log_prob=True, - input_shape=Shape.STATES, - hiddens=[256, 256], - hidden_activation=["relu", "relu"], - output_shape=Shape.ACTIONS, - output_activation=None) + model = categorical_model( + observation_space=observation_space, + action_space=action_space, + device=device, + unnormalized_log_prob=True, + input_shape=Shape.STATES, + hiddens=[256, 256], + hidden_activation=["relu", "relu"], + output_shape=Shape.ACTIONS, + output_activation=None, + ) model.to(device=device) observations = torch.ones((10, model.num_observations), device=device) output = model.act({"states": observations}) assert output[0].shape == (10, 1) -@hypothesis.given(observation_space_size=st.integers(min_value=1, max_value=10), - action_space_size=st.integers(min_value=1, max_value=10)) + +@hypothesis.given( + observation_space_size=st.integers(min_value=1, max_value=10), + action_space_size=st.integers(min_value=1, max_value=10), +) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_deterministic_model(capsys, observation_space_size, action_space_size, device): observation_space = gym.spaces.Box(np.array([-1] * observation_space_size), np.array([1] * observation_space_size)) action_space = gym.spaces.Box(np.array([-1] * action_space_size), np.array([1] * action_space_size)) # TODO: randomize all parameters - model = deterministic_model(observation_space=observation_space, - action_space=action_space, - device=device, - clip_actions=False, - input_shape=Shape.STATES, - hiddens=[256, 256], - hidden_activation=["relu", "relu"], - output_shape=Shape.ACTIONS, - output_activation=None, - output_scale=1) + model = deterministic_model( + observation_space=observation_space, + action_space=action_space, + device=device, + clip_actions=False, + input_shape=Shape.STATES, + hiddens=[256, 256], + hidden_activation=["relu", "relu"], + output_shape=Shape.ACTIONS, + output_activation=None, + output_scale=1, + ) model.to(device=device) observations = torch.ones((10, model.num_observations), device=device) output = model.act({"states": observations}) assert output[0].shape == (10, model.num_actions) -@hypothesis.given(observation_space_size=st.integers(min_value=1, max_value=10), - action_space_size=st.integers(min_value=1, max_value=10)) + +@hypothesis.given( + observation_space_size=st.integers(min_value=1, max_value=10), + action_space_size=st.integers(min_value=1, max_value=10), +) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_gaussian_model(capsys, observation_space_size, action_space_size, device): observation_space = gym.spaces.Box(np.array([-1] * observation_space_size), np.array([1] * observation_space_size)) action_space = gym.spaces.Box(np.array([-1] * action_space_size), np.array([1] * action_space_size)) # TODO: randomize all parameters - model = gaussian_model(observation_space=observation_space, - action_space=action_space, - device=device, - clip_actions=False, - clip_log_std=True, - min_log_std=-20, - max_log_std=2, - initial_log_std=0, - input_shape=Shape.STATES, - hiddens=[256, 256], - hidden_activation=["relu", "relu"], - output_shape=Shape.ACTIONS, - output_activation=None, - output_scale=1) + model = gaussian_model( + observation_space=observation_space, + action_space=action_space, + device=device, + clip_actions=False, + clip_log_std=True, + min_log_std=-20, + max_log_std=2, + initial_log_std=0, + input_shape=Shape.STATES, + hiddens=[256, 256], + hidden_activation=["relu", "relu"], + output_shape=Shape.ACTIONS, + output_activation=None, + output_scale=1, + ) model.to(device=device) observations = torch.ones((10, model.num_observations), device=device) output = model.act({"states": observations}) assert output[0].shape == (10, model.num_actions) -@hypothesis.given(observation_space_size=st.integers(min_value=1, max_value=10), - action_space_size=st.integers(min_value=1, max_value=10)) + +@hypothesis.given( + observation_space_size=st.integers(min_value=1, max_value=10), + action_space_size=st.integers(min_value=1, max_value=10), +) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_multivariate_gaussian_model(capsys, observation_space_size, action_space_size, device): observation_space = gym.spaces.Box(np.array([-1] * observation_space_size), np.array([1] * observation_space_size)) action_space = gym.spaces.Box(np.array([-1] * action_space_size), np.array([1] * action_space_size)) # TODO: randomize all parameters - model = multivariate_gaussian_model(observation_space=observation_space, - action_space=action_space, - device=device, - clip_actions=False, - clip_log_std=True, - min_log_std=-20, - max_log_std=2, - initial_log_std=0, - input_shape=Shape.STATES, - hiddens=[256, 256], - hidden_activation=["relu", "relu"], - output_shape=Shape.ACTIONS, - output_activation=None, - output_scale=1) + model = multivariate_gaussian_model( + observation_space=observation_space, + action_space=action_space, + device=device, + clip_actions=False, + clip_log_std=True, + min_log_std=-20, + max_log_std=2, + initial_log_std=0, + input_shape=Shape.STATES, + hiddens=[256, 256], + hidden_activation=["relu", "relu"], + output_shape=Shape.ACTIONS, + output_activation=None, + output_scale=1, + ) model.to(device=device) observations = torch.ones((10, model.num_observations), device=device) output = model.act({"states": observations}) assert output[0].shape == (10, model.num_actions) -@hypothesis.given(observation_space_size=st.integers(min_value=1, max_value=10), - action_space_size=st.integers(min_value=1, max_value=10)) + +@hypothesis.given( + observation_space_size=st.integers(min_value=1, max_value=10), + action_space_size=st.integers(min_value=1, max_value=10), +) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"]) def test_shared_model(capsys, observation_space_size, action_space_size, device): observation_space = gym.spaces.Box(np.array([-1] * observation_space_size), np.array([1] * observation_space_size)) action_space = gym.spaces.Box(np.array([-1] * action_space_size), np.array([1] * action_space_size)) # TODO: randomize all parameters - model = shared_model(observation_space=observation_space, - action_space=action_space, - device=device, - structure="", - roles=["policy", "value"], - parameters=[ - { - "clip_actions": False, - "clip_log_std": True, - "min_log_std": -20, - "max_log_std": 2, - "initial_log_std": 0, - "input_shape": Shape.STATES, - "hiddens": [256, 256], - "hidden_activation": ["relu", "relu"], - "output_shape": Shape.ACTIONS, - "output_activation": None, - "output_scale": 1, - }, - { - "clip_actions": False, - "input_shape": Shape.STATES, - "hiddens": [256, 256], - "hidden_activation": ["relu", "relu"], - "output_shape": Shape.ONE, - "output_activation": None, - "output_scale": 1, - } - ], - single_forward_pass=True) + model = shared_model( + observation_space=observation_space, + action_space=action_space, + device=device, + structure="", + roles=["policy", "value"], + parameters=[ + { + "clip_actions": False, + "clip_log_std": True, + "min_log_std": -20, + "max_log_std": 2, + "initial_log_std": 0, + "input_shape": Shape.STATES, + "hiddens": [256, 256], + "hidden_activation": ["relu", "relu"], + "output_shape": Shape.ACTIONS, + "output_activation": None, + "output_scale": 1, + }, + { + "clip_actions": False, + "input_shape": Shape.STATES, + "hiddens": [256, 256], + "hidden_activation": ["relu", "relu"], + "output_shape": Shape.ONE, + "output_activation": None, + "output_scale": 1, + }, + ], + single_forward_pass=True, + ) model.to(device=device) observations = torch.ones((10, model.num_observations), device=device) diff --git a/tests/torch/test_torch_model_instantiators_definition.py b/tests/torch/test_torch_model_instantiators_definition.py index b4c9d52a..796e000d 100644 --- a/tests/torch/test_torch_model_instantiators_definition.py +++ b/tests/torch/test_torch_model_instantiators_definition.py @@ -14,7 +14,7 @@ deterministic_model, gaussian_model, multivariate_gaussian_model, - shared_model + shared_model, ) from skrl.utils.model_instantiators.torch.common import _generate_modules, _get_activation_function, _parse_input @@ -31,6 +31,7 @@ def test_get_activation_function(capsys): assert activation is not None, f"{item} -> None" exec(activation, _globals, {}) + def test_parse_input(capsys): # check for Shape enum (compatibility with prior versions) for input in [Shape.STATES, Shape.OBSERVATIONS, Shape.ACTIONS, Shape.STATES_ACTIONS]: @@ -52,6 +53,7 @@ def test_parse_input(capsys): output = _parse_input(str(input)) assert output.replace("'", '"') == statement, f"'{output}' != '{statement}'" + def test_generate_modules(capsys): _globals = {"nn": torch.nn} @@ -147,6 +149,7 @@ def test_generate_modules(capsys): assert isinstance(container, torch.nn.Sequential) assert len(container) == 2 + def test_gaussian_model(capsys): device = "cpu" observation_space = gym.spaces.Box(np.array([-1] * 5), np.array([1] * 5)) @@ -170,19 +173,15 @@ def test_gaussian_model(capsys): """ content = yaml.safe_load(content) # source - model = gaussian_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=True, - **content) + model = gaussian_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=True, **content + ) with capsys.disabled(): print(model) # instance - model = gaussian_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=False, - **content) + model = gaussian_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=False, **content + ) model.to(device=device) with capsys.disabled(): print(model) @@ -191,6 +190,7 @@ def test_gaussian_model(capsys): output = model.act({"states": observations}) assert output[0].shape == (10, 2) + def test_multivariate_gaussian_model(capsys): device = "cpu" observation_space = gym.spaces.Box(np.array([-1] * 5), np.array([1] * 5)) @@ -214,19 +214,15 @@ def test_multivariate_gaussian_model(capsys): """ content = yaml.safe_load(content) # source - model = multivariate_gaussian_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=True, - **content) + model = multivariate_gaussian_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=True, **content + ) with capsys.disabled(): print(model) # instance - model = multivariate_gaussian_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=False, - **content) + model = multivariate_gaussian_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=False, **content + ) model.to(device=device) with capsys.disabled(): print(model) @@ -235,6 +231,7 @@ def test_multivariate_gaussian_model(capsys): output = model.act({"states": observations}) assert output[0].shape == (10, 2) + def test_deterministic_model(capsys): device = "cpu" observation_space = gym.spaces.Box(np.array([-1] * 5), np.array([1] * 5)) @@ -255,19 +252,15 @@ def test_deterministic_model(capsys): """ content = yaml.safe_load(content) # source - model = deterministic_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=True, - **content) + model = deterministic_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=True, **content + ) with capsys.disabled(): print(model) # instance - model = deterministic_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=False, - **content) + model = deterministic_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=False, **content + ) model.to(device=device) with capsys.disabled(): print(model) @@ -276,6 +269,7 @@ def test_deterministic_model(capsys): output = model.act({"states": observations}) assert output[0].shape == (10, 3) + def test_categorical_model(capsys): device = "cpu" observation_space = gym.spaces.Box(np.array([-1] * 5), np.array([1] * 5)) @@ -295,19 +289,15 @@ def test_categorical_model(capsys): """ content = yaml.safe_load(content) # source - model = categorical_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=True, - **content) + model = categorical_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=True, **content + ) with capsys.disabled(): print(model) # instance - model = categorical_model(observation_space=observation_space, - action_space=action_space, - device=device, - return_source=False, - **content) + model = categorical_model( + observation_space=observation_space, action_space=action_space, device=device, return_source=False, **content + ) model.to(device=device) with capsys.disabled(): print(model) @@ -316,6 +306,7 @@ def test_categorical_model(capsys): output = model.act({"states": observations}) assert output[0].shape == (10, 1) + def test_shared_model(capsys): device = "cpu" observation_space = gym.spaces.Box(np.array([-1] * 5), np.array([1] * 5)) @@ -352,27 +343,31 @@ def test_shared_model(capsys): content_policy = yaml.safe_load(content_policy) content_value = yaml.safe_load(content_value) # source - model = shared_model(observation_space=observation_space, - action_space=action_space, - device=device, - roles=["policy", "value"], - parameters=[ - content_policy, - content_value, - ], - return_source=True) + model = shared_model( + observation_space=observation_space, + action_space=action_space, + device=device, + roles=["policy", "value"], + parameters=[ + content_policy, + content_value, + ], + return_source=True, + ) with capsys.disabled(): print(model) # instance - model = shared_model(observation_space=observation_space, - action_space=action_space, - device=device, - roles=["policy", "value"], - parameters=[ - content_policy, - content_value, - ], - return_source=False) + model = shared_model( + observation_space=observation_space, + action_space=action_space, + device=device, + roles=["policy", "value"], + parameters=[ + content_policy, + content_value, + ], + return_source=False, + ) model.to(device=device) with capsys.disabled(): print(model) diff --git a/tests/torch/test_torch_utils_spaces.py b/tests/torch/test_torch_utils_spaces.py index aa3a08ac..bc9c8126 100644 --- a/tests/torch/test_torch_utils_spaces.py +++ b/tests/torch/test_torch_utils_spaces.py @@ -14,7 +14,7 @@ sample_space, tensorize_space, unflatten_tensorized_space, - untensorize_space + untensorize_space, ) from ..stategies import gym_space_stategy, gymnasium_space_stategy @@ -28,6 +28,7 @@ def _check_backend(x, backend): else: raise ValueError(f"Invalid backend type: {backend}") + def check_sampled_space(space, x, n, backend): if isinstance(space, gymnasium.spaces.Box): _check_backend(x, backend) @@ -65,6 +66,7 @@ def occupied_size(s): space_size = compute_space_size(space, occupied_size=True) assert space_size == occupied_size(space) + @hypothesis.given(space=gymnasium_space_stategy()) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_tensorize_space(capsys, space: gymnasium.spaces.Space): @@ -96,6 +98,7 @@ def check_tensorized_space(s, x, n): tensorized_space = tensorize_space(space, sampled_space) check_tensorized_space(space, tensorized_space, 5) + @hypothesis.given(space=gymnasium_space_stategy()) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_untensorize_space(capsys, space: gymnasium.spaces.Space): @@ -107,7 +110,9 @@ def check_untensorized_space(s, x, squeeze_batch_dimension): assert isinstance(x, (np.ndarray, int)) assert isinstance(x, int) if squeeze_batch_dimension else x.shape == (1, 1) elif isinstance(s, gymnasium.spaces.MultiDiscrete): - assert isinstance(x, np.ndarray) and x.shape == s.nvec.shape if squeeze_batch_dimension else (1, *s.nvec.shape) + assert ( + isinstance(x, np.ndarray) and x.shape == s.nvec.shape if squeeze_batch_dimension else (1, *s.nvec.shape) + ) elif isinstance(s, gymnasium.spaces.Dict): list(map(check_untensorized_space, s.values(), x.values(), [squeeze_batch_dimension] * len(s))) elif isinstance(s, gymnasium.spaces.Tuple): @@ -123,6 +128,7 @@ def check_untensorized_space(s, x, squeeze_batch_dimension): untensorized_space = untensorize_space(space, tensorized_space, squeeze_batch_dimension=True) check_untensorized_space(space, untensorized_space, squeeze_batch_dimension=True) + @hypothesis.given(space=gymnasium_space_stategy(), batch_size=st.integers(min_value=1, max_value=10)) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_sample_space(capsys, space: gymnasium.spaces.Space, batch_size: int): @@ -133,6 +139,7 @@ def test_sample_space(capsys, space: gymnasium.spaces.Space, batch_size: int): sampled_space = sample_space(space, batch_size, backend="torch") check_sampled_space(space, sampled_space, batch_size, backend="torch") + @hypothesis.given(space=gymnasium_space_stategy()) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_flatten_tensorized_space(capsys, space: gymnasium.spaces.Space): @@ -146,6 +153,7 @@ def test_flatten_tensorized_space(capsys, space: gymnasium.spaces.Space): flattened_space = flatten_tensorized_space(tensorized_space) assert flattened_space.shape == (5, space_size) + @hypothesis.given(space=gymnasium_space_stategy()) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_unflatten_tensorized_space(capsys, space: gymnasium.spaces.Space): @@ -159,6 +167,7 @@ def test_unflatten_tensorized_space(capsys, space: gymnasium.spaces.Space): unflattened_space = unflatten_tensorized_space(space, flattened_space) check_sampled_space(space, unflattened_space, 5, backend="torch") + @hypothesis.given(space=gym_space_stategy()) @hypothesis.settings(suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], deadline=None) def test_convert_gym_space(capsys, space: gym.spaces.Space): diff --git a/tests/torch/test_torch_wrapper_deepmind.py b/tests/torch/test_torch_wrapper_deepmind.py index 57c46ca4..4d880e52 100644 --- a/tests/torch/test_torch_wrapper_deepmind.py +++ b/tests/torch/test_torch_wrapper_deepmind.py @@ -28,7 +28,10 @@ def test_env(capsys: pytest.CaptureFixture): # check properties assert env.state_space is None - assert isinstance(env.observation_space, gym.Space) and sorted(list(env.observation_space.keys())) == ["orientation", "velocity"] + assert isinstance(env.observation_space, gym.Space) and sorted(list(env.observation_space.keys())) == [ + "orientation", + "velocity", + ] assert isinstance(env.action_space, gym.Space) and env.action_space.shape == (1,) assert isinstance(env.num_envs, int) and env.num_envs == num_envs assert isinstance(env.num_agents, int) and env.num_agents == 1 diff --git a/tests/torch/test_torch_wrapper_gym.py b/tests/torch/test_torch_wrapper_gym.py index cfee7672..440688b2 100644 --- a/tests/torch/test_torch_wrapper_gym.py +++ b/tests/torch/test_torch_wrapper_gym.py @@ -46,6 +46,7 @@ def test_env(capsys: pytest.CaptureFixture): env.close() + @pytest.mark.parametrize("vectorization_mode", ["async", "sync"]) def test_vectorized_env(capsys: pytest.CaptureFixture, vectorization_mode: str): num_envs = 10 diff --git a/tests/torch/test_torch_wrapper_gymnasium.py b/tests/torch/test_torch_wrapper_gymnasium.py index f603801e..80bbc154 100644 --- a/tests/torch/test_torch_wrapper_gymnasium.py +++ b/tests/torch/test_torch_wrapper_gymnasium.py @@ -45,6 +45,7 @@ def test_env(capsys: pytest.CaptureFixture): env.close() + @pytest.mark.parametrize("vectorization_mode", ["async", "sync"]) def test_vectorized_env(capsys: pytest.CaptureFixture, vectorization_mode: str): num_envs = 10 diff --git a/tests/torch/test_torch_wrapper_isaacgym.py b/tests/torch/test_torch_wrapper_isaacgym.py index e7d2b367..c8614270 100644 --- a/tests/torch/test_torch_wrapper_isaacgym.py +++ b/tests/torch/test_torch_wrapper_isaacgym.py @@ -27,7 +27,7 @@ def __init__(self, num_states) -> None: self.state_space = gym.spaces.Box(np.ones(self.num_states) * -np.Inf, np.ones(self.num_states) * np.Inf) self.observation_space = gym.spaces.Box(np.ones(self.num_obs) * -np.Inf, np.ones(self.num_obs) * np.Inf) - self.action_space = gym.spaces.Box(np.ones(self.num_actions) * -1., np.ones(self.num_actions) * 1.) + self.action_space = gym.spaces.Box(np.ones(self.num_actions) * -1.0, np.ones(self.num_actions) * 1.0) def reset(self) -> Dict[str, torch.Tensor]: obs_dict = {} diff --git a/tests/torch/test_torch_wrapper_isaaclab.py b/tests/torch/test_torch_wrapper_isaaclab.py index 66fffcbe..dca49e39 100644 --- a/tests/torch/test_torch_wrapper_isaaclab.py +++ b/tests/torch/test_torch_wrapper_isaaclab.py @@ -23,6 +23,7 @@ Dict[AgentID, dict], ] + class IsaacLabEnv(gym.Env): def __init__(self, num_states) -> None: self.num_actions = 1 @@ -58,7 +59,9 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) def step(self, action: torch.Tensor) -> VecEnvStepReturn: assert action.clone().shape == torch.Size([self.num_envs, 1]) - observations = {"policy": torch.ones((self.num_envs, self.num_observations), device=self.device, dtype=torch.float32)} + observations = { + "policy": torch.ones((self.num_envs, self.num_observations), device=self.device, dtype=torch.float32) + } rewards = torch.zeros(self.num_envs, device=self.device, dtype=torch.float32) terminated = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) truncated = torch.zeros_like(terminated) @@ -99,9 +102,7 @@ def _configure_env_spaces(self): if not self.num_states: self.state_space = None if self.num_states < 0: - self.state_space = gym.spaces.Box( - low=-np.inf, high=np.inf, shape=(sum(self.num_observations.values()),) - ) + self.state_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(sum(self.num_observations.values()),)) else: self.state_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.num_states,)) @@ -109,16 +110,28 @@ def _configure_env_spaces(self): def unwrapped(self): return self - def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Dict[AgentID, ObsType], dict]: - observations = {agent: torch.ones((self.num_envs, self.num_observations[agent]), device=self.device) for agent in self.possible_agents} + def reset( + self, seed: int | None = None, options: dict[str, Any] | None = None + ) -> tuple[Dict[AgentID, ObsType], dict]: + observations = { + agent: torch.ones((self.num_envs, self.num_observations[agent]), device=self.device) + for agent in self.possible_agents + } return observations, self.extras def step(self, action: Dict[AgentID, torch.Tensor]) -> EnvStepReturn: for agent in self.possible_agents: assert action[agent].clone().shape == torch.Size([self.num_envs, self.num_actions[agent]]) - observations = {agent: torch.ones((self.num_envs, self.num_observations[agent]), device=self.device) for agent in self.possible_agents} - rewards = {agent: torch.zeros(self.num_envs, device=self.device, dtype=torch.float32) for agent in self.possible_agents} - terminated = {agent: torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) for agent in self.possible_agents} + observations = { + agent: torch.ones((self.num_envs, self.num_observations[agent]), device=self.device) + for agent in self.possible_agents + } + rewards = { + agent: torch.zeros(self.num_envs, device=self.device, dtype=torch.float32) for agent in self.possible_agents + } + terminated = { + agent: torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) for agent in self.possible_agents + } truncated = {agent: torch.zeros_like(terminated[agent]) for agent in self.possible_agents} return observations, rewards, terminated, truncated, self.extras @@ -176,6 +189,7 @@ def test_env(capsys: pytest.CaptureFixture, num_states: int): env.close() + @pytest.mark.parametrize("num_states", [0, 5]) def test_multi_agent_env(capsys: pytest.CaptureFixture, num_states: int): num_envs = 10 @@ -212,7 +226,9 @@ def test_multi_agent_env(capsys: pytest.CaptureFixture, num_states: int): assert isinstance(observation, Mapping) assert isinstance(info, Mapping) for i, agent in enumerate(possible_agents): - assert isinstance(observation[agent], torch.Tensor) and observation[agent].shape == torch.Size([num_envs, i + 20]) + assert isinstance(observation[agent], torch.Tensor) and observation[agent].shape == torch.Size( + [num_envs, i + 20] + ) for _ in range(3): observation, reward, terminated, truncated, info = env.step(action) state = env.state() @@ -223,10 +239,16 @@ def test_multi_agent_env(capsys: pytest.CaptureFixture, num_states: int): assert isinstance(truncated, Mapping) assert isinstance(info, Mapping) for i, agent in enumerate(possible_agents): - assert isinstance(observation[agent], torch.Tensor) and observation[agent].shape == torch.Size([num_envs, i + 20]) + assert isinstance(observation[agent], torch.Tensor) and observation[agent].shape == torch.Size( + [num_envs, i + 20] + ) assert isinstance(reward[agent], torch.Tensor) and reward[agent].shape == torch.Size([num_envs, 1]) - assert isinstance(terminated[agent], torch.Tensor) and terminated[agent].shape == torch.Size([num_envs, 1]) - assert isinstance(truncated[agent], torch.Tensor) and truncated[agent].shape == torch.Size([num_envs, 1]) + assert isinstance(terminated[agent], torch.Tensor) and terminated[agent].shape == torch.Size( + [num_envs, 1] + ) + assert isinstance(truncated[agent], torch.Tensor) and truncated[agent].shape == torch.Size( + [num_envs, 1] + ) if num_states: assert isinstance(state, torch.Tensor) and state.shape == torch.Size([num_envs, num_states]) else: diff --git a/tests/torch/test_torch_wrapper_omniisaacgym.py b/tests/torch/test_torch_wrapper_omniisaacgym.py index 8d0e899e..5525e947 100644 --- a/tests/torch/test_torch_wrapper_omniisaacgym.py +++ b/tests/torch/test_torch_wrapper_omniisaacgym.py @@ -25,9 +25,16 @@ def __init__(self, num_states) -> None: self.device = "cpu" # initialize data spaces (defaults to gym.Box) - self.action_space = gym.spaces.Box(np.ones(self.num_actions, dtype=np.float32) * -1.0, np.ones(self.num_actions, dtype=np.float32) * 1.0) - self.observation_space = gym.spaces.Box(np.ones(self.num_observations, dtype=np.float32) * -np.Inf, np.ones(self.num_observations, dtype=np.float32) * np.Inf) - self.state_space = gym.spaces.Box(np.ones(self.num_states, dtype=np.float32) * -np.Inf, np.ones(self.num_states, dtype=np.float32) * np.Inf) + self.action_space = gym.spaces.Box( + np.ones(self.num_actions, dtype=np.float32) * -1.0, np.ones(self.num_actions, dtype=np.float32) * 1.0 + ) + self.observation_space = gym.spaces.Box( + np.ones(self.num_observations, dtype=np.float32) * -np.Inf, + np.ones(self.num_observations, dtype=np.float32) * np.Inf, + ) + self.state_space = gym.spaces.Box( + np.ones(self.num_states, dtype=np.float32) * -np.Inf, np.ones(self.num_states, dtype=np.float32) * np.Inf + ) def reset(self): observations = {"obs": torch.ones((self.num_envs, self.num_observations), device=self.device)} @@ -35,7 +42,9 @@ def reset(self): def step(self, actions): assert actions.clone().shape == torch.Size([self.num_envs, 1]) - observations = {"obs": torch.ones((self.num_envs, self.num_observations), device=self.device, dtype=torch.float32)} + observations = { + "obs": torch.ones((self.num_envs, self.num_observations), device=self.device, dtype=torch.float32) + } rewards = torch.zeros(self.num_envs, device=self.device, dtype=torch.float32) dones = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) return observations, rewards, dones, self.extras diff --git a/tests/torch/test_torch_wrapper_pettingzoo.py b/tests/torch/test_torch_wrapper_pettingzoo.py index d6d1d860..ebe0531e 100644 --- a/tests/torch/test_torch_wrapper_pettingzoo.py +++ b/tests/torch/test_torch_wrapper_pettingzoo.py @@ -31,7 +31,11 @@ def test_env(capsys: pytest.CaptureFixture): assert isinstance(env.action_spaces, Mapping) and len(env.action_spaces) == num_agents for agent in possible_agents: assert isinstance(env.state_space(agent), gym.Space) and env.state_space(agent).shape == (560, 880, 3) - assert isinstance(env.observation_space(agent), gym.Space) and env.observation_space(agent).shape == (457, 120, 3) + assert isinstance(env.observation_space(agent), gym.Space) and env.observation_space(agent).shape == ( + 457, + 120, + 3, + ) assert isinstance(env.action_space(agent), gym.Space) and env.action_space(agent).shape == (1,) assert isinstance(env.possible_agents, list) and sorted(env.possible_agents) == sorted(possible_agents) assert isinstance(env.num_envs, int) and env.num_envs == num_envs @@ -47,7 +51,9 @@ def test_env(capsys: pytest.CaptureFixture): assert isinstance(observation, Mapping) assert isinstance(info, Mapping) for agent in possible_agents: - assert isinstance(observation[agent], torch.Tensor) and observation[agent].shape == torch.Size([num_envs, math.prod((457, 120, 3))]) + assert isinstance(observation[agent], torch.Tensor) and observation[agent].shape == torch.Size( + [num_envs, math.prod((457, 120, 3))] + ) for _ in range(3): observation, reward, terminated, truncated, info = env.step(action) state = env.state() @@ -58,10 +64,16 @@ def test_env(capsys: pytest.CaptureFixture): assert isinstance(truncated, Mapping) assert isinstance(info, Mapping) for agent in possible_agents: - assert isinstance(observation[agent], torch.Tensor) and observation[agent].shape == torch.Size([num_envs, math.prod((457, 120, 3))]) + assert isinstance(observation[agent], torch.Tensor) and observation[agent].shape == torch.Size( + [num_envs, math.prod((457, 120, 3))] + ) assert isinstance(reward[agent], torch.Tensor) and reward[agent].shape == torch.Size([num_envs, 1]) - assert isinstance(terminated[agent], torch.Tensor) and terminated[agent].shape == torch.Size([num_envs, 1]) - assert isinstance(truncated[agent], torch.Tensor) and truncated[agent].shape == torch.Size([num_envs, 1]) + assert isinstance(terminated[agent], torch.Tensor) and terminated[agent].shape == torch.Size( + [num_envs, 1] + ) + assert isinstance(truncated[agent], torch.Tensor) and truncated[agent].shape == torch.Size( + [num_envs, 1] + ) assert isinstance(state, torch.Tensor) and state.shape == torch.Size([num_envs, math.prod((560, 880, 3))]) env.close() diff --git a/tests/utils.py b/tests/utils.py index e6df75f0..9c3b6ef2 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -5,7 +5,7 @@ class DummyEnv(gym.Env): - def __init__(self, num_envs, device = "cpu"): + def __init__(self, num_envs, device="cpu"): self.num_agents = 1 self.num_envs = num_envs self.device = torch.device(device) @@ -46,7 +46,9 @@ class _DummyBaseAgent: def __init__(self): pass - def record_transition(self, states, actions, rewards, next_states, terminated, truncated, infos, timestep, timesteps): + def record_transition( + self, states, actions, rewards, next_states, terminated, truncated, infos, timestep, timesteps + ): pass def pre_interaction(self, timestep, timesteps): @@ -69,7 +71,9 @@ def init(self, trainer_cfg=None): def act(self, states, timestep, timesteps): return torch.tensor([]), None, {} - def record_transition(self, states, actions, rewards, next_states, terminated, truncated, infos, timestep, timesteps): + def record_transition( + self, states, actions, rewards, next_states, terminated, truncated, infos, timestep, timesteps + ): pass def pre_interaction(self, timestep, timesteps):