diff --git a/colab/Colab_UnityEnvironment_4_SB3VectorEnv.ipynb b/colab/Colab_UnityEnvironment_4_SB3VectorEnv.ipynb index 84bb94c52c5..a93d638830f 100644 --- a/colab/Colab_UnityEnvironment_4_SB3VectorEnv.ipynb +++ b/colab/Colab_UnityEnvironment_4_SB3VectorEnv.ipynb @@ -161,8 +161,8 @@ "from pathlib import Path\n", "from typing import Callable, Any\n", "\n", - "import gym\n", - "from gym import Env\n", + "import gymnasium as gym\n", + "from gymnasium import Env\n", "\n", "from stable_baselines3 import PPO\n", "from stable_baselines3.common.vec_env import VecMonitor, VecEnv, SubprocVecEnv\n", diff --git a/com.unity.ml-agents/CHANGELOG.md b/com.unity.ml-agents/CHANGELOG.md index fa6ac19a6be..26102b8958e 100755 --- a/com.unity.ml-agents/CHANGELOG.md +++ b/com.unity.ml-agents/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Major Changes +#### ml-agents / ml-agents-envs +- Migrated from gym to gymnasium (#6309) + ### Minor Changes #### com.unity.ml-agents (C#) - Fixed StackingSensor compressed observation for sensors with more than 3 channels. (#6299) diff --git a/com.unity.ml-agents/Documentation~/ML-Agents-Overview.md b/com.unity.ml-agents/Documentation~/ML-Agents-Overview.md index 126b1148e59..b84a816231f 100644 --- a/com.unity.ml-agents/Documentation~/ML-Agents-Overview.md +++ b/com.unity.ml-agents/Documentation~/ML-Agents-Overview.md @@ -28,8 +28,8 @@ The ML-Agents Toolkit contains five high-level components: - **Python Low-Level API** - which contains a low-level Python interface for interacting and manipulating a learning environment. Note that, unlike the Learning Environment, the Python API is not part of Unity, but lives outside and communicates with Unity through the Communicator. This API is contained in a dedicated `mlagents_envs` Python package and is used by the Python training process to communicate with and control the Academy during training. However, it can be used for other purposes as well. For example, you could use the API to use Unity as the simulation engine for your own machine learning algorithms. See [Python API](Python-LLAPI.md) for more information. - **External Communicator** - which connects the Learning Environment with the Python Low-Level API. It lives within the Learning Environment. - **Python Trainers** which contains all the machine learning algorithms that enable training agents. The algorithms are implemented in Python and are part of their own `mlagents` Python package. The package exposes a single command-line utility `mlagents-learn` that supports all the training methods and options outlined in this document. The Python Trainers interface solely with the Python Low-Level API. -- **Gym Wrapper** (not pictured). A common way in which machine learning researchers interact with simulation environments is via a wrapper provided by OpenAI called [gym](https://github.com/openai/gym). We provide a gym wrapper in the `ml-agents-envs` package and [instructions](Python-Gym-API.md) for using it with existing machine learning algorithms which utilize gym. -- **PettingZoo Wrapper** (not pictured) PettingZoo is python API for interacting with multi-agent simulation environments that provides a gym-like interface. We provide a PettingZoo wrapper for Unity ML-Agents environments in the `ml-agents-envs` package and [instructions](Python-PettingZoo-API.md) for using it with machine learning algorithms. +- **Gym Wrapper** (not pictured). A common way in which machine learning researchers interact with simulation environments is via a wrapper provided by the Farama Foundation called [gymnasium](https://gymnasium.farama.org/) (formerly OpenAI `gym`). Unity provides a gym wrapper in the `ml-agents-envs` package and [instructions](Python-Gym-API.md) for using it with existing machine learning algorithms which utilize gymnasium. +- **PettingZoo Wrapper** (not pictured) PettingZoo is python API for interacting with multi-agent simulation environments that provides a gym-like interface. Unity provides a PettingZoo wrapper for Unity ML-Agents environments in the `ml-agents-envs` package and [instructions](Python-PettingZoo-API.md) for using it with machine learning algorithms.

Simplified ML-Agents Scene Block Diagram

@@ -68,7 +68,7 @@ It is important to note that the ML-Agents Toolkit leverages [Sentis](Inference- ### Custom Training and Inference -In the previous mode, the Agents were used for training to generate a PyTorch model that the Agents can later use. However, any user of the ML-Agents Toolkit can leverage their own algorithms for training. In this case, the behaviors of all the Agents in the scene will be controlled within Python. You can even turn your environment into a [gym.](Python-Gym-API.md) +In the previous mode, the Agents were used for training to generate a PyTorch model that the Agents can later use. However, any user of the ML-Agents Toolkit can leverage their own algorithms for training. In this case, the behaviors of all the Agents in the scene will be controlled within Python. You can even turn your environment into a [gymnasium environment.](Python-Gym-API.md) Unity doesn't provide a tutorial highlighting this mode, but you can learn more about the Python API in [Unity ML-Agents Python Low Level API](Python-LLAPI.md). diff --git a/com.unity.ml-agents/Documentation~/Python-Gym-API-Documentation.md b/com.unity.ml-agents/Documentation~/Python-Gym-API-Documentation.md index f3db6d59cb8..0c9eeb1cf90 100644 --- a/com.unity.ml-agents/Documentation~/Python-Gym-API-Documentation.md +++ b/com.unity.ml-agents/Documentation~/Python-Gym-API-Documentation.md @@ -42,19 +42,19 @@ Environment initialization #### reset ```python - | reset() -> Union[List[np.ndarray], np.ndarray] + | reset(*, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None) -> Tuple[np.ndarray, Dict] ``` -Resets the state of the environment and returns an initial observation. Returns: observation (object/list): the initial observation of the space. +Resets the state of the environment and returns an initial observation and info. Returns: observation (object/list): the initial observation of the space. info (dict): contains auxiliary diagnostic information. #### step ```python - | step(action: List[Any]) -> GymStepResult + | step(action: Any) -> GymStepResult ``` -Run one timestep of the environment's dynamics. When end of episode is reached, you are responsible for calling `reset()` to reset this environment's state. Accepts an action and returns a tuple (observation, reward, done, info). +Run one timestep of the environment's dynamics. When end of episode is reached, you are responsible for calling `reset()` to reset this environment's state. Accepts an action and returns a tuple (observation, reward, terminated, truncated, info). **Arguments**: @@ -63,14 +63,15 @@ Run one timestep of the environment's dynamics. When end of episode is reached, **Returns**: - `observation` _object/list_ - agent's observation of the current environment reward (float/list) : amount of reward returned after previous action -- `done` _boolean/list_ - whether the episode has ended. +- `terminated` _boolean/list_ - whether the episode has ended by termination. +- `truncated` _boolean/list_ - whether the episode has ended by truncation. - `info` _dict_ - contains auxiliary diagnostic information. #### render ```python - | render(mode="rgb_array") + | render() ``` Return the latest visual observations. Note that it will not render a new frame of the environment. diff --git a/com.unity.ml-agents/Documentation~/Python-Gym-API.md b/com.unity.ml-agents/Documentation~/Python-Gym-API.md index bf7d835daab..6b3c597d86b 100644 --- a/com.unity.ml-agents/Documentation~/Python-Gym-API.md +++ b/com.unity.ml-agents/Documentation~/Python-Gym-API.md @@ -1,8 +1,8 @@ # Unity ML-Agents Gym Wrapper -A common way in which machine learning researchers interact with simulation environments is via a wrapper provided by OpenAI called `gym`. For more information on the gym interface, see [here](https://github.com/openai/gym). +A common way in which machine learning researchers interact with simulation environments is via a wrapper provided by the Farama Foundation called `gymnasium` (formerly known as OpenAI `gym`). For more information on the gymnasium interface, visit the [Gymnasium Documentation](https://gymnasium.farama.org/index.html), and the [Gymnasium Github repository](https://github.com/Farama-Foundation/Gymnasium). -We provide a gym wrapper and instructions for using it with existing machine learning algorithms which utilize gym. Our wrapper provides interfaces on top of our `UnityEnvironment` class, which is the default way of interfacing with a Unity environment via Python. +The ML Agents package provides a gym wrapper and instructions for using it with existing machine learning algorithms which utilize gymnasium. This wrapper provides interfaces on top of the `UnityEnvironment` class, which is the default way of interfacing with a Unity environment via Python. ## Installation @@ -11,7 +11,7 @@ The gym wrapper is part of the `mlagents_envs` package. Please refer to the [mla ## Using the Gym Wrapper -The gym interface is available from `gym_unity.envs`. To launch an environment from the root of the project repository use: +The gym interface is available from `mlagents_envs.envs.unity_gym_env`. To launch an environment from the root of the project repository use: ```python from mlagents_envs.envs.unity_gym_env import UnityToGymWrapper @@ -23,13 +23,13 @@ env = UnityToGymWrapper(unity_env, uint8_visual, flatten_branched, allow_multipl - `uint8_visual` refers to whether to output visual observations as `uint8` values (0-255). Many common Gym environments (e.g. Atari) do this. By default, they will be floats (0.0-1.0). Defaults to `False`. -- `flatten_branched` will flatten a branched discrete action space into a Gym Discrete. Otherwise, it will be converted into a MultiDiscrete. Defaults to `False`. +- `flatten_branched` will flatten a branched discrete action space into a gymnasium `Discrete` space. Otherwise, it will be converted into a `MultiDiscrete`. Defaults to `False`. - `allow_multiple_obs` will return a list of observations. The first elements contain the visual observations and the last element contains the array of vector observations. If False the environment returns a single array (containing a single visual observations, if present, otherwise the vector observation). Defaults to `False`. - `action_space_seed` is the optional seed for action sampling. If non-None, will be used to set the random seed on created gym.Space instances. -The returned environment `env` will function as a gym. +The returned environment `env` will function as a gymnasium environment. ## Limitations @@ -40,130 +40,121 @@ The returned environment `env` will function as a gym. - Environment registration for use with `gym.make()` is currently not supported. - Calling env.render() will not render a new frame of the environment. It will return the latest visual observation if using visual observations. -## Running OpenAI Baselines Algorithms +## Training with Stable-Baselines3 -OpenAI provides a set of open-source maintained and tested Reinforcement Learning algorithms called the [Baselines](https://github.com/openai/baselines). +[Stable-Baselines3](https://github.com/DLR-RM/stable-baselines3) (SB3) is a set of reliable, actively maintained implementations of reinforcement learning algorithms in PyTorch. It is the community successor to OpenAI Baselines and, like this wrapper, is built on the Farama Foundation `gymnasium` API, so you can train ML-Agents environments with it directly. -Using the provided Gym wrapper, it is possible to train ML-Agents environments using these algorithms. This requires the creation of custom training scripts to launch each algorithm. In most cases these scripts can be created by making slight modifications to the ones provided for Atari and Mujoco environments. +Install SB3 with: -These examples were tested with baselines version 0.1.6. - -### Example - DQN Baseline +```sh +pip install stable-baselines3 +``` -To train an agent to play the `GridWorld` environment using the Baselines DQN algorithm, you first need to install the Baselines package. For instructions, refer to the [Baselines README](https://github.com/openai/baselines). +### Example - PPO -Next, create a file called `train_unity.py`. Then create an `/envs/` directory and build the environment to that directory. For more information on building Unity environments, see [here](Learning-Environment-Executable.md). Note that because of limitations of the DQN baseline, the environment must have a single visual observation, a single discrete action and a single Agent in the scene. Add the following code to the `train_unity.py` file: +To train an agent with PPO on a single-agent environment, create a file called `train_unity.py` with the following code. Then create an `/envs/` directory and build the environment to that directory. For more information on building Unity environments, refer to [Using an Environment Executable](Learning-Environment-Executable.md). ```python -import gym - -from baselines import deepq -from baselines import logger +from stable_baselines3 import PPO from mlagents_envs.environment import UnityEnvironment from mlagents_envs.envs.unity_gym_env import UnityToGymWrapper def main(): - unity_env = UnityEnvironment( < path - to - environment >) - env = UnityToGymWrapper(unity_env, uint8_visual=True) - logger.configure('./logs') # Change to log in a different directory - act = deepq.learn( - env, - "cnn", # For visual inputs - lr=2.5e-4, - total_timesteps=1000000, - buffer_size=50000, - exploration_fraction=0.05, - exploration_final_eps=0.1, - print_freq=20, - train_freq=5, - learning_starts=20000, - target_network_update_freq=50, - gamma=0.99, - prioritized_replay=False, - checkpoint_freq=1000, - checkpoint_path='./logs', # Change to save model in a different directory - dueling=True - ) - print("Saving model to unity_model.pkl") - act.save("unity_model.pkl") - - -if __name__ == '__main__': - main() + unity_env = UnityEnvironment() + env = UnityToGymWrapper(unity_env) + model = PPO("MlpPolicy", env, verbose=1) + model.learn(total_timesteps=100000) + print("Saving model to unity_model.zip") + model.save("unity_model") + + +if __name__ == "__main__": + main() ``` -To start the training process, run the following from the directory containing -`train_unity.py`: +To start the training process, run the following from the directory containing `train_unity.py`: ```sh -python -m train_unity +python train_unity.py ``` -### Other Algorithms - -Other algorithms in the Baselines repository can be run using scripts similar to the examples from the baselines package. In most cases, the primary changes needed to use a Unity environment are to import `UnityToGymWrapper`, and to replace the environment creation code, typically `gym.make()`, with a call to `UnityToGymWrapper(unity_environment)` passing the environment as input. +Use the `MlpPolicy` for environments with vector observations, and the `CnnPolicy` for environments with visual (image) observations. -A typical rule of thumb is that for vision-based environments, modification should be done to Atari training scripts, and for vector observation environments, modification should be done to Mujoco scripts. +### Example - DQN -Some algorithms will make use of `make_env()` or `make_mujoco_env()` functions. You can define a similar function for Unity environments. An example of such a method using the PPO2 baseline: +DQN requires a discrete action space. For an environment with a single visual observation and branched discrete actions, enable `uint8_visual` (to output `uint8` image observations) and `flatten_branched` (to collapse the branched action space into a single `Discrete` space): ```python +from stable_baselines3 import DQN + from mlagents_envs.environment import UnityEnvironment -from mlagents_envs.envs import UnityToGymWrapper -from baselines.common.vec_env.subproc_vec_env import SubprocVecEnv -from baselines.common.vec_env.dummy_vec_env import DummyVecEnv -from baselines.bench import Monitor -from baselines import logger -import baselines.ppo2.ppo2 as ppo2 +from mlagents_envs.envs.unity_gym_env import UnityToGymWrapper + -import os +def main(): + unity_env = UnityEnvironment() + env = UnityToGymWrapper(unity_env, uint8_visual=True, flatten_branched=True) + model = DQN( + "CnnPolicy", + env, + learning_rate=2.5e-4, + buffer_size=50000, + learning_starts=20000, + target_update_interval=50, + gamma=0.99, + verbose=1, + ) + model.learn(total_timesteps=1000000) + print("Saving model to unity_model.zip") + model.save("unity_model") + + +if __name__ == "__main__": + main() +``` -try: - from mpi4py import MPI -except ImportError: - MPI = None +### Training on multiple environments in parallel +SB3 can train on several environment instances at once using a vectorized environment. Each Unity instance must use a distinct `base_port` so the instances don't conflict: -def make_unity_env(env_directory, num_env, visual, start_index=0): - """ - Create a wrapped, monitored Unity environment. - """ +```python +from stable_baselines3 import PPO +from stable_baselines3.common.vec_env import SubprocVecEnv - def make_env(rank, use_visual=True): # pylint: disable=C0111 - def _thunk(): - unity_env = UnityEnvironment(env_directory, base_port=5000 + rank) - env = UnityToGymWrapper(unity_env, uint8_visual=True) - env = Monitor(env, logger.get_dir() and os.path.join(logger.get_dir(), str(rank))) - return env +from mlagents_envs.environment import UnityEnvironment +from mlagents_envs.envs.unity_gym_env import UnityToGymWrapper - return _thunk - if visual: +def make_unity_env(env_directory, num_env, start_index=0): + """Create a wrapped, vectorized Unity environment.""" + + def make_env(rank): + def _thunk(): + unity_env = UnityEnvironment(env_directory, base_port=5000 + rank) + return UnityToGymWrapper(unity_env) + + return _thunk + return SubprocVecEnv([make_env(i + start_index) for i in range(num_env)]) - else: - rank = MPI.COMM_WORLD.Get_rank() if MPI else 0 - return DummyVecEnv([make_env(rank, use_visual=False)]) def main(): - env = make_unity_env( < path - to - environment >, 4, True) - ppo2.learn( - network="mlp", - env=env, - total_timesteps=100000, - lr=1e-3, - ) + env = make_unity_env(, 4) + model = PPO("MlpPolicy", env, verbose=1) + model.learn(total_timesteps=100000) -if __name__ == '__main__': - main() +if __name__ == "__main__": + main() ``` ## Run Google Dopamine Algorithms -Google provides a framework [Dopamine](https://github.com/google/dopamine), and implementations of algorithms, e.g. DQN, Rainbow, and the C51 variant of Rainbow. Using the Gym wrapper, we can run Unity environments using Dopamine. +> [!NOTE] The following walkthrough was written for an older, OpenAI `gym`-based release of Dopamine. This wrapper now follows the Farama Foundation `gymnasium` API: `reset()` returns `(observation, info)` and `step()` returns `(observation, reward, terminated, truncated, info)`. Recent versions of Dopamine target `gymnasium` and are compatible with the wrapper, but the exact file names, module paths, and configuration steps described here may differ from the version you install. Treat this section as a general guide rather than a step-by-step recipe. + +Google provides a framework [Dopamine](https://github.com/google/dopamine), and implementations of algorithms, e.g. DQN, Rainbow, and the C51 variant of Rainbow. Using the Gym wrapper, you can run Unity environments using Dopamine. First, after installing the Gym wrapper, clone the Dopamine repository. @@ -177,11 +168,11 @@ Then, follow the appropriate install instructions as specified on [Dopamine's ho First, open `dopamine/atari/run_experiment.py`. Alternatively, copy the entire `atari` folder, and name it something else (e.g. `unity`). If you choose the copy approach, be sure to change the package names in the import statements in `train.py` to your new directory. -Within `run_experiment.py`, we will need to make changes to which environment is instantiated, just as in the Baselines example. At the top of the file, insert +Within `run_experiment.py`, you need to make changes to which environment is instantiated, just as in the Stable-Baselines3 examples in the previous section. At the top of the file, insert ```python from mlagents_envs.environment import UnityEnvironment -from mlagents_envs.envs import UnityToGymWrapper +from mlagents_envs.envs.unity_gym_env import UnityToGymWrapper ``` to import the Gym Wrapper. Navigate to the `create_atari_environment` method in the same file, and switch to instantiating a Unity environment by replacing the method with the following code. @@ -194,13 +185,13 @@ to import the Gym Wrapper. Navigate to the `create_atari_environment` method in return env ``` -`` is the path to your built Unity executable. For more information on building Unity environments, see [here](Learning-Environment-Executable.md), and note the Limitations section below. +`` is the path to your built Unity executable. For more information on building Unity environments, refer to [Using an Environment Executable](Learning-Environment-Executable.md), and note the Limitations section below. -Note that we are not using the preprocessor from Dopamine, as it uses many Atari-specific calls. Furthermore, frame-skipping can be done from within Unity, rather than on the Python side. +Note that the script is not using the preprocessor from Dopamine, as it uses many Atari-specific calls. Furthermore, frame-skipping can be done from within Unity, rather than on the Python side. ### Limitations -Since Dopamine is designed around variants of DQN, it is only compatible with discrete action spaces, and specifically the Discrete Gym space. For environments that use branched discrete action spaces, you can enable the `flatten_branched` parameter in `UnityToGymWrapper`, which treats each combination of branched actions as separate actions. +Since Dopamine is designed around variants of DQN, it is only compatible with discrete action spaces, and specifically the `Discrete` gymnasium space. For environments that use branched discrete action spaces, you can enable the `flatten_branched` parameter in `UnityToGymWrapper`, which treats each combination of branched actions as separate actions. Furthermore, when building your environments, ensure that your Agent is using visual observations with greyscale enabled, and that the dimensions of the visual observations is 84 by 84 (matches the parameter found in `dqn_agent.py` and `rainbow_agent.py`). Dopamine's agents currently do not automatically adapt to the observation dimensions or number of channels. @@ -256,7 +247,7 @@ python -um dopamine.unity.train \ --gin_files='dopamine/agents/rainbow/configs/rainbow.gin' ``` -Again, we assume that you've copied `atari` into a separate folder. Remember to replace `unity` with the directory you copied your files into. If you edited the Atari files directly, this should be `atari`. +Again, ensure to copy `atari` into a separate folder. Remember to replace `unity` with the directory you copied your files into. If you edited the Atari files directly, this should be `atari`. ### Example: GridWorld diff --git a/com.unity.ml-agents/Documentation~/Python-PettingZoo-API-Documentation.md b/com.unity.ml-agents/Documentation~/Python-PettingZoo-API-Documentation.md index 5040ac6879a..d8db8cc484d 100644 --- a/com.unity.ml-agents/Documentation~/Python-PettingZoo-API-Documentation.md +++ b/com.unity.ml-agents/Documentation~/Python-PettingZoo-API-Documentation.md @@ -57,7 +57,7 @@ Initializes a Unity AEC environment wrapper. | step(action: Any) -> None ``` -Sets the action of the active agent and get the observation, reward, done and info of the next agent. +Sets the action of the active agent and get the observation, reward, termination, truncation and info of the next agent. **Arguments**: @@ -79,7 +79,7 @@ Returns the observation an agent currently can make. `last()` calls this functio | last(observe=True) ``` -returns observation, cumulative reward, done, info for the current agent (specified by self.agent_selection) +returns observation, cumulative reward, termination, truncation, info for the current agent (specified by self.agent_selection) # mlagents\_envs.envs.unity\_parallel\_env @@ -180,31 +180,23 @@ The side channels of the environment. You can access the side channels of an env #### reset ```python - | reset() + | reset(seed: Optional[int] = None, options: Optional[Dict] = None) -> Any ``` Resets the environment. - -#### seed - -```python - | seed(seed=None) -``` - -Reseeds the environment (making the resulting environment deterministic). -`reset()` must be called after `seed()`, and before `step()`. - #### render ```python - | render(mode="human") + | render() ``` NOT SUPPORTED. -Displays a rendered frame from the environment, if supported. Alternate render modes in the default environments are `'rgb_array'` which returns a numpy array and is supported by all environments outside of classic, and `'ansi'` which returns the strings printed (specific to classic environments). +Renders the environment as specified by self.render_mode, if supported. + +Render mode can be `human` to display a window. Other render modes in the default environments are `'rgb_array'` which returns a numpy array and is supported by all environments outside of classic, and `'ansi'` which returns the strings printed (specific to classic environments). #### close diff --git a/com.unity.ml-agents/Documentation~/Python-PettingZoo-API.md b/com.unity.ml-agents/Documentation~/Python-PettingZoo-API.md index dd3f5d26d8e..b77ca9db566 100644 --- a/com.unity.ml-agents/Documentation~/Python-PettingZoo-API.md +++ b/com.unity.ml-agents/Documentation~/Python-PettingZoo-API.md @@ -16,13 +16,13 @@ This wrapper is compatible with PettingZoo API. Please check out [PettingZoo API ```python from mlagents_envs.environment import UnityEnvironment -from mlagents_envs.envs import UnityToPettingZooWrapper +from mlagents_envs.envs.unity_aec_env import UnityAECEnv unity_env = UnityEnvironment("StrikersVsGoalie") -env = UnityToPettingZooWrapper(unity_env) +env = UnityAECEnv(unity_env) env.reset() for agent in env.agent_iter(): - observation, reward, done, info = env.last() + observation, reward, termination, truncation, info = env.last() action = policy(observation, agent) env.step(action) ``` diff --git a/ml-agents-envs/colabs/Colab_PettingZoo.ipynb b/ml-agents-envs/colabs/Colab_PettingZoo.ipynb index 0cc4c8bbdda..f975d5aaad0 100644 --- a/ml-agents-envs/colabs/Colab_PettingZoo.ipynb +++ b/ml-agents-envs/colabs/Colab_PettingZoo.ipynb @@ -207,10 +207,10 @@ "\n", "env.reset()\n", "for agent in env.agent_iter(env.num_agents * num_cycles):\n", - " prev_observe, reward, done, info = env.last()\n", + " prev_observe, reward, termination, truncation, info = env.last()\n", " if isinstance(prev_observe, dict) and 'action_mask' in prev_observe:\n", " action_mask = prev_observe['action_mask']\n", - " if done:\n", + " if termination or truncation:\n", " action = None\n", " else:\n", " action = env.action_spaces[agent].sample() # randomly choose an action for example\n", diff --git a/ml-agents-envs/mlagents_envs/envs/env_helpers.py b/ml-agents-envs/mlagents_envs/envs/env_helpers.py index 768e6706038..c4c42410d6e 100644 --- a/ml-agents-envs/mlagents_envs/envs/env_helpers.py +++ b/ml-agents-envs/mlagents_envs/envs/env_helpers.py @@ -17,7 +17,8 @@ def _unwrap_batch_steps(batch_steps, behavior_name): termination_id = [ _behavior_to_agent_id(behavior_name, i) for i in termination_batch.agent_id ] - agents = decision_id + termination_id + agents = list(dict.fromkeys(decision_id + termination_id)) + obs = { agent_id: [batch_obs[i] for batch_obs in termination_batch.obs] for i, agent_id in enumerate(termination_id) @@ -40,30 +41,46 @@ def _unwrap_batch_steps(batch_steps, behavior_name): } ) obs = {k: v if len(v) > 1 else v[0] for k, v in obs.items()} - dones = {agent_id: True for agent_id in termination_id} - dones.update({agent_id: False for agent_id in decision_id}) rewards = { - agent_id: termination_batch.reward[i] - for i, agent_id in enumerate(termination_id) + agent_id: decision_batch.reward[i] for i, agent_id in enumerate(decision_id) } rewards.update( - {agent_id: decision_batch.reward[i] for i, agent_id in enumerate(decision_id)} + { + agent_id: termination_batch.reward[i] + for i, agent_id in enumerate(termination_id) + } ) cumulative_rewards = {k: v for k, v in rewards.items()} infos = {} + terminations = {} + truncations = {} for i, agent_id in enumerate(decision_id): infos[agent_id] = {} infos[agent_id]["behavior_name"] = behavior_name infos[agent_id]["group_id"] = decision_batch.group_id[i] infos[agent_id]["group_reward"] = decision_batch.group_reward[i] + truncations[agent_id] = False + terminations[agent_id] = False for i, agent_id in enumerate(termination_id): infos[agent_id] = {} infos[agent_id]["behavior_name"] = behavior_name infos[agent_id]["group_id"] = termination_batch.group_id[i] infos[agent_id]["group_reward"] = termination_batch.group_reward[i] - infos[agent_id]["interrupted"] = termination_batch.interrupted[i] + truncated = bool(termination_batch.interrupted[i]) + infos[agent_id]["interrupted"] = truncated + truncations[agent_id] = truncated + terminations[agent_id] = not truncated id_map = {agent_id: i for i, agent_id in enumerate(decision_id)} - return agents, obs, dones, rewards, cumulative_rewards, infos, id_map + return ( + agents, + obs, + terminations, + truncations, + rewards, + cumulative_rewards, + infos, + id_map, + ) def _parse_behavior(full_behavior): diff --git a/ml-agents-envs/mlagents_envs/envs/unity_aec_env.py b/ml-agents-envs/mlagents_envs/envs/unity_aec_env.py index 4bb6fdf3909..72e9fc350ba 100644 --- a/ml-agents-envs/mlagents_envs/envs/unity_aec_env.py +++ b/ml-agents-envs/mlagents_envs/envs/unity_aec_env.py @@ -1,5 +1,5 @@ from typing import Any, Optional -from gym import error +from gymnasium import error from mlagents_envs.base_env import BaseEnv from pettingzoo import AECEnv @@ -22,8 +22,8 @@ def __init__(self, env: BaseEnv, seed: Optional[int] = None): def step(self, action: Any) -> None: """ - Sets the action of the active agent and get the observation, reward, done - and info of the next agent. + Sets the action of the active agent and get the observation, reward, + termination, truncation and info of the next agent. :param action: The action for the active agent """ self._assert_loaded() @@ -53,16 +53,20 @@ def observe(self, agent_id): return ( self._observations[agent_id], self._cumm_rewards[agent_id], - self._dones[agent_id], + self._terminations[agent_id], + self._truncations[agent_id], self._infos[agent_id], ) def last(self, observe=True): """ - returns observation, cumulative reward, done, info for the current agent (specified by self.agent_selection) + returns observation, cumulative reward, termination, truncation, info + for the current agent (specified by self.agent_selection) """ - obs, reward, done, info = self.observe(self._agents[self._agent_index]) - return obs if observe else None, reward, done, info + obs, cumm_rewards, terminated, truncated, info = self.observe( + self._agents[self._agent_index] + ) + return obs if observe else None, cumm_rewards, terminated, truncated, info @property def agent_selection(self): diff --git a/ml-agents-envs/mlagents_envs/envs/unity_gym_env.py b/ml-agents-envs/mlagents_envs/envs/unity_gym_env.py index df29a95c9ab..1a052cdb4fc 100644 --- a/ml-agents-envs/mlagents_envs/envs/unity_gym_env.py +++ b/ml-agents-envs/mlagents_envs/envs/unity_gym_env.py @@ -3,8 +3,8 @@ import numpy as np from typing import Any, Dict, List, Optional, Tuple, Union -import gym -from gym import error, spaces +import gymnasium as gym +from gymnasium import error, spaces from mlagents_envs.base_env import ActionTuple, BaseEnv from mlagents_envs.base_env import DecisionSteps, TerminalSteps @@ -20,7 +20,7 @@ class UnityGymException(error.Error): logger = logging_util.get_logger(__name__) -GymStepResult = Tuple[np.ndarray, float, bool, Dict] +GymStepResult = Tuple[np.ndarray, float, bool, bool, Dict] class UnityToGymWrapper(gym.Env): @@ -151,11 +151,16 @@ def __init__( else: self._observation_space = list_spaces[0] # only return the first one - def reset(self) -> Union[List[np.ndarray], np.ndarray]: - """Resets the state of the environment and returns an initial observation. - Returns: observation (object/list): the initial observation of the + def reset( + self, *, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None + ) -> Tuple[np.ndarray, Dict]: + """Resets the state of the environment and returns an initial observation and info. + Returns: + observation (object/list): the initial observation of the space. + info (dict): contains auxiliary diagnostic information. """ + super().reset(seed=seed, options=options) self._env.reset() decision_step, _ = self._env.get_steps(self.name) n_agents = len(decision_step) @@ -163,26 +168,27 @@ def reset(self) -> Union[List[np.ndarray], np.ndarray]: self.game_over = False res: GymStepResult = self._single_step(decision_step) - return res[0] + return res[0], res[4] - def step(self, action: List[Any]) -> GymStepResult: + def step(self, action: Any) -> GymStepResult: """Run one timestep of the environment's dynamics. When end of episode is reached, you are responsible for calling `reset()` to reset this environment's state. - Accepts an action and returns a tuple (observation, reward, done, info). + Accepts an action and returns a tuple (observation, reward, terminated, truncated, info). Args: action (object/list): an action provided by the environment Returns: observation (object/list): agent's observation of the current environment reward (float/list) : amount of reward returned after previous action - done (boolean/list): whether the episode has ended. + terminated (boolean/list): whether the episode has ended by termination. + truncated (boolean/list): whether the episode has ended by truncation. info (dict): contains auxiliary diagnostic information. """ if self.game_over: raise UnityGymException( "You are calling 'step()' even though this environment has already " - "returned done = True. You must always call 'reset()' once you " - "receive 'done = True'." + "returned `terminated` or `truncated` as True. You must always call 'reset()' once you " + "receive `terminated` or `truncated` as True." ) if self._flattener is not None: # Translate action into list @@ -227,9 +233,19 @@ def _single_step(self, info: Union[DecisionSteps, TerminalSteps]) -> GymStepResu visual_obs = self._get_vis_obs_list(info) self.visual_obs = self._preprocess_single(visual_obs[0][0]) - done = isinstance(info, TerminalSteps) + if isinstance(info, TerminalSteps): + interrupted = info.interrupted[0] + terminated, truncated = not interrupted, interrupted + else: + terminated, truncated = False, False - return (default_observation, info.reward[0], done, {"step": info}) + return ( + default_observation, + info.reward[0], + terminated, + truncated, + {"step": info}, + ) def _preprocess_single(self, single_visual_obs: np.ndarray) -> np.ndarray: if self.uint8_visual: @@ -276,7 +292,7 @@ def _get_vec_obs_size(self) -> int: result += obs_spec.shape[0] return result - def render(self, mode="rgb_array"): + def render(self): """ Return the latest visual observations. Note that it will not render a new frame of the environment. @@ -306,7 +322,7 @@ def _check_agents(n_agents: int) -> None: @property def metadata(self): - return {"render.modes": ["rgb_array"]} + return {"render_modes": ["rgb_array"]} @property def reward_range(self) -> Tuple[float, float]: diff --git a/ml-agents-envs/mlagents_envs/envs/unity_parallel_env.py b/ml-agents-envs/mlagents_envs/envs/unity_parallel_env.py index 09398d27fa8..9dcce070a88 100644 --- a/ml-agents-envs/mlagents_envs/envs/unity_parallel_env.py +++ b/ml-agents-envs/mlagents_envs/envs/unity_parallel_env.py @@ -1,5 +1,5 @@ from typing import Optional, Dict, Any, Tuple -from gym import error +from gymnasium import error from mlagents_envs.base_env import BaseEnv from pettingzoo import ParallelEnv @@ -20,13 +20,17 @@ def __init__(self, env: BaseEnv, seed: Optional[int] = None): """ super().__init__(env, seed) - def reset(self) -> Dict[str, Any]: + def reset( + self, + seed: Optional[int] = None, + options: Optional[Dict] = None, + ) -> Tuple[Dict[str, Any], Dict[str, Dict]]: """ Resets the environment. """ - super().reset() + super().reset(seed=seed, options=options) - return self._observations + return self._observations, self._infos def step(self, actions: Dict[str, Any]) -> Tuple: self._assert_loaded() @@ -50,4 +54,10 @@ def step(self, actions: Dict[str, Any]) -> Tuple: self._cleanup_agents() self._live_agents.sort() # unnecessary, only for passing API test - return self._observations, self._rewards, self._dones, self._infos + return ( + self._observations, + self._rewards, + self._terminations, + self._truncations, + self._infos, + ) diff --git a/ml-agents-envs/mlagents_envs/envs/unity_pettingzoo_base_env.py b/ml-agents-envs/mlagents_envs/envs/unity_pettingzoo_base_env.py index 3457f18c882..36d75b7ba87 100644 --- a/ml-agents-envs/mlagents_envs/envs/unity_pettingzoo_base_env.py +++ b/ml-agents-envs/mlagents_envs/envs/unity_pettingzoo_base_env.py @@ -1,7 +1,7 @@ import atexit from typing import Optional, List, Set, Dict, Any, Tuple import numpy as np -from gym import error, spaces +from gymnasium import error, spaces from mlagents_envs.base_env import BaseEnv, ActionTuple from mlagents_envs.envs.env_helpers import _agent_id_to_behavior, _unwrap_batch_steps @@ -17,7 +17,9 @@ def __init__( super().__init__() atexit.register(self.close) self._env = env - self.metadata = metadata + self.metadata = ( + metadata if metadata is not None else {"name": "unity_ml-agents"} + ) self._assert_loaded() self._agent_index = 0 @@ -32,7 +34,8 @@ def __init__( self._possible_agents: Set[str] = set() # all agents that have ever appear self._agent_id_to_index: Dict[str, int] = {} # agent_id: index in decision step self._observations: Dict[str, np.ndarray] = {} # agent_id: obs - self._dones: Dict[str, bool] = {} # agent_id: done + self._terminations: Dict[str, bool] = {} # agent_id: terminated + self._truncations: Dict[str, bool] = {} # agent_id: truncated self._rewards: Dict[str, float] = {} # agent_id: reward self._cumm_rewards: Dict[str, float] = {} # agent_id: reward self._infos: Dict[str, Dict] = {} # agent_id: info @@ -45,7 +48,7 @@ def __init__( if not self._env.behavior_specs: self._env.step() for behavior_name in self._env.behavior_specs.keys(): - _, _, _ = self._batch_update(behavior_name) + _, _, _, _ = self._batch_update(behavior_name) self._update_observation_spaces() self._update_action_spaces() @@ -132,7 +135,7 @@ def _update_action_spaces(self) -> None: continue if act_spec.continuous_size > 0: c_space = spaces.Box( - -1, 1, (act_spec.continuous_size,), dtype=np.int32 + -1, 1, (act_spec.continuous_size,), dtype=np.float32 ) if self._seed is not None: c_space.seed(self._seed) @@ -162,7 +165,7 @@ def _process_action(self, current_agent, action): else: action = ActionTuple(action, None) - if not self._dones[current_agent]: + if not (self._terminations[current_agent] or self._truncations[current_agent]): current_behavior = _agent_id_to_behavior(current_agent) current_index = self._agent_id_to_index[current_agent] if action.continuous is not None: @@ -176,7 +179,8 @@ def _process_action(self, current_agent, action): else: self._live_agents.remove(current_agent) del self._observations[current_agent] - del self._dones[current_agent] + del self._terminations[current_agent] + del self._truncations[current_agent] del self._rewards[current_agent] del self._cumm_rewards[current_agent] del self._infos[current_agent] @@ -187,15 +191,18 @@ def _step(self): self._env.step() self._reset_states() for behavior_name in self._env.behavior_specs.keys(): - dones, rewards, cumulative_rewards = self._batch_update(behavior_name) - self._dones.update(dones) + terminations, truncations, rewards, cumulative_rewards = self._batch_update( + behavior_name + ) + self._terminations.update(terminations) + self._truncations.update(truncations) self._rewards.update(rewards) self._cumm_rewards.update(cumulative_rewards) self._agent_index = 0 def _cleanup_agents(self): - for current_agent, done in self.dones.items(): - if done: + for current_agent in self._agents: + if self._terminations[current_agent] or self._truncations[current_agent]: self._live_agents.remove(current_agent) @property @@ -226,25 +233,38 @@ def _reset_states(self): self._live_agents = [] self._agents = [] self._observations = {} - self._dones = {} + self._terminations = {} + self._truncations = {} self._rewards = {} self._cumm_rewards = {} self._infos = {} self._agent_id_to_index = {} - def reset(self): + def reset( + self, + seed: Optional[int] = None, + options: Optional[Dict] = None, + ) -> Any: """ Resets the environment. """ + self._seed = seed + if seed is not None: + # Action spaces are created (and seeded) once in __init__, so reseed + # the existing ones here to make reset(seed=...) actually reproducible. + for space in self._action_spaces.values(): + space.seed(seed) + self._assert_loaded() self._agent_index = 0 self._reset_states() self._possible_agents = set() self._env.reset() for behavior_name in self._env.behavior_specs.keys(): - _, _, _ = self._batch_update(behavior_name) + _, _, _, _ = self._batch_update(behavior_name) self._live_agents.sort() # unnecessary, only for passing API test - self._dones = {agent: False for agent in self._agents} + self._terminations = {agent: False for agent in self._agents} + self._truncations = {agent: False for agent in self._agents} self._rewards = {agent: 0 for agent in self._agents} self._cumm_rewards = {agent: 0 for agent in self._agents} @@ -256,7 +276,8 @@ def _batch_update(self, behavior_name): ( agents, obs, - dones, + terminations, + truncations, rewards, cumulative_rewards, infos, @@ -268,29 +289,28 @@ def _batch_update(self, behavior_name): self._infos.update(infos) self._agent_id_to_index.update(id_map) self._possible_agents.update(agents) - return dones, rewards, cumulative_rewards - - def seed(self, seed=None): - """ - Reseeds the environment (making the resulting environment deterministic). - `reset()` must be called after `seed()`, and before `step()`. - """ - self._seed = seed + return terminations, truncations, rewards, cumulative_rewards - def render(self, mode="human"): + def render(self): """ NOT SUPPORTED. - Displays a rendered frame from the environment, if supported. - Alternate render modes in the default environments are `'rgb_array'` + Renders the environment as specified by self.render_mode, if supported. + + Render mode can be `human` to display a window. + Other render modes in the default environments are `'rgb_array'` which returns a numpy array and is supported by all environments outside of classic, and `'ansi'` which returns the strings printed (specific to classic environments). """ pass @property - def dones(self): - return dict(self._dones) + def terminations(self): + return dict(self._terminations) + + @property + def truncations(self): + return dict(self._truncations) @property def agents(self): diff --git a/ml-agents-envs/setup.py b/ml-agents-envs/setup.py index 5e4c80b0863..2c55a95cd00 100644 --- a/ml-agents-envs/setup.py +++ b/ml-agents-envs/setup.py @@ -58,9 +58,9 @@ def run(self): "Pillow>=4.2.1", "protobuf>=3.6,<3.21", "pyyaml>=3.1.0", - "gym>=0.21.0", - "pettingzoo==1.15.0", - "numpy>=1.23.5,<1.24.0", + "gymnasium>=0.26.0", + "pettingzoo>=1.24.0", + "numpy>=1.23.5,<2.0", "filelock>=3.4.0", ], python_requires=">=3.10.1,<=3.10.12", diff --git a/ml-agents-envs/tests/test_gym.py b/ml-agents-envs/tests/test_gym.py index 4fc2bf548cc..b3e4e1a0d84 100644 --- a/ml-agents-envs/tests/test_gym.py +++ b/ml-agents-envs/tests/test_gym.py @@ -2,7 +2,7 @@ import pytest import numpy as np -from gym import spaces +from gymnasium import spaces from mlagents_envs.envs.unity_gym_env import UnityToGymWrapper from mlagents_envs.base_env import ( @@ -23,17 +23,42 @@ def test_gym_wrapper(): mock_env, mock_spec, mock_decision_step, mock_terminal_step ) env = UnityToGymWrapper(mock_env) - assert isinstance(env.reset(), np.ndarray) + reset_obs, reset_info = env.reset() + assert isinstance(reset_obs, np.ndarray) + assert isinstance(reset_info, dict) actions = env.action_space.sample() assert actions.shape[0] == 2 - obs, rew, done, info = env.step(actions) + obs, rew, term, trunc, info = env.step(actions) assert env.observation_space.contains(obs) assert isinstance(obs, np.ndarray) assert isinstance(rew, float) - assert isinstance(done, (bool, np.bool_)) + assert isinstance(term, (bool, np.bool_)) + assert isinstance(trunc, (bool, np.bool_)) assert isinstance(info, dict) +@pytest.mark.parametrize( + "interrupted,expected", + [(True, (False, True)), (False, (True, False))], + ids=["truncated", "terminated"], +) +def test_gym_wrapper_terminated_truncated(interrupted, expected): + # A non-empty TerminalSteps drives step() down its terminal branch so the + # interrupted -> (terminated, truncated) translation is actually exercised. + mock_env = mock.MagicMock() + mock_spec = create_mock_group_spec() + mock_decision_step, _ = create_mock_vector_steps(mock_spec) + mock_terminal_step = create_mock_terminal_steps(mock_spec, interrupted=interrupted) + setup_mock_unityenvironment( + mock_env, mock_spec, mock_decision_step, mock_terminal_step + ) + env = UnityToGymWrapper(mock_env) + env.reset() + obs, rew, term, trunc, info = env.step(env.action_space.sample()) + assert (term, trunc) == expected + assert env.game_over is True + + def test_branched_flatten(): mock_env = mock.MagicMock() mock_spec = create_mock_group_spec( @@ -108,14 +133,17 @@ def test_gym_wrapper_visual(use_uint8): env = UnityToGymWrapper(mock_env, uint8_visual=use_uint8) assert isinstance(env.observation_space, spaces.Box) - assert isinstance(env.reset(), np.ndarray) + reset_obs, reset_info = env.reset() + assert isinstance(reset_obs, np.ndarray) + assert isinstance(reset_info, dict) actions = env.action_space.sample() assert actions.shape[0] == 2 - obs, rew, done, info = env.step(actions) + obs, rew, term, trunc, info = env.step(actions) assert env.observation_space.contains(obs) assert isinstance(obs, np.ndarray) assert isinstance(rew, float) - assert isinstance(done, (bool, np.bool_)) + assert isinstance(term, (bool, np.bool_)) + assert isinstance(trunc, (bool, np.bool_)) assert isinstance(info, dict) @@ -137,32 +165,35 @@ def test_gym_wrapper_single_visual_and_vector(use_uint8): env = UnityToGymWrapper(mock_env, uint8_visual=use_uint8, allow_multiple_obs=True) assert isinstance(env.observation_space, spaces.Tuple) assert len(env.observation_space) == 2 - reset_obs = env.reset() + reset_obs, reset_info = env.reset() assert isinstance(reset_obs, list) + assert isinstance(reset_info, dict) assert len(reset_obs) == 2 assert all(isinstance(ob, np.ndarray) for ob in reset_obs) assert reset_obs[-1].shape == (3,) assert len(reset_obs[0].shape) == 3 actions = env.action_space.sample() assert actions.shape == (2,) - obs, rew, done, info = env.step(actions) + obs, rew, term, trunc, info = env.step(actions) assert isinstance(obs, list) assert len(obs) == 2 assert all(isinstance(ob, np.ndarray) for ob in obs) assert reset_obs[-1].shape == (3,) assert isinstance(rew, float) - assert isinstance(done, (bool, np.bool_)) + assert isinstance(term, (bool, np.bool_)) + assert isinstance(trunc, (bool, np.bool_)) assert isinstance(info, dict) # check behavior for allow_multiple_obs = False env = UnityToGymWrapper(mock_env, uint8_visual=use_uint8, allow_multiple_obs=False) assert isinstance(env.observation_space, spaces.Box) - reset_obs = env.reset() + reset_obs, reset_info = env.reset() assert isinstance(reset_obs, np.ndarray) + assert isinstance(reset_info, dict) assert len(reset_obs.shape) == 3 actions = env.action_space.sample() assert actions.shape == (2,) - obs, rew, done, info = env.step(actions) + obs, rew, term, trunc, info = env.step(actions) assert isinstance(obs, np.ndarray) @@ -184,28 +215,31 @@ def test_gym_wrapper_multi_visual_and_vector(use_uint8): env = UnityToGymWrapper(mock_env, uint8_visual=use_uint8, allow_multiple_obs=True) assert isinstance(env.observation_space, spaces.Tuple) assert len(env.observation_space) == 3 - reset_obs = env.reset() + reset_obs, reset_info = env.reset() assert isinstance(reset_obs, list) + assert isinstance(reset_info, dict) assert len(reset_obs) == 3 assert all(isinstance(ob, np.ndarray) for ob in reset_obs) assert reset_obs[-1].shape == (3,) actions = env.action_space.sample() assert actions.shape == (2,) - obs, rew, done, info = env.step(actions) + obs, rew, term, trunc, info = env.step(actions) assert all(isinstance(ob, np.ndarray) for ob in obs) assert isinstance(rew, float) - assert isinstance(done, (bool, np.bool_)) + assert isinstance(term, (bool, np.bool_)) + assert isinstance(trunc, (bool, np.bool_)) assert isinstance(info, dict) # check behavior for allow_multiple_obs = False env = UnityToGymWrapper(mock_env, uint8_visual=use_uint8, allow_multiple_obs=False) assert isinstance(env.observation_space, spaces.Box) - reset_obs = env.reset() + reset_obs, reset_info = env.reset() assert isinstance(reset_obs, np.ndarray) + assert isinstance(reset_info, dict) assert len(reset_obs.shape) == 3 actions = env.action_space.sample() assert actions.shape == (2,) - obs, rew, done, info = env.step(actions) + obs, rew, term, trunc, info = env.step(actions) assert isinstance(obs, np.ndarray) @@ -264,6 +298,20 @@ def create_mock_vector_steps(specs, num_agents=1, number_visual_observations=0): ) +def create_mock_terminal_steps(specs, num_agents=1, interrupted=True): + """ + Creates a non-empty TerminalSteps mock so the wrapper's terminal branch + (and the interrupted -> terminated/truncated split) can be tested. + """ + obs = [np.array(num_agents * [[1, 2, 3]], dtype=np.float32).reshape(num_agents, 3)] + rewards = np.array(num_agents * [1.0]) + interrupted_arr = np.array(num_agents * [interrupted], dtype=bool) + agents = np.array(range(0, num_agents)) + group_id = np.array(num_agents * [0]) + group_rewards = np.array(num_agents * [0.0]) + return TerminalSteps(obs, rewards, interrupted_arr, agents, group_id, group_rewards) + + def setup_mock_unityenvironment(mock_env, mock_spec, mock_decision, mock_termination): """ Takes a mock UnityEnvironment and adds the appropriate properties, defined by the mock diff --git a/ml-agents-envs/tests/test_pettingzoo_wrapper.py b/ml-agents-envs/tests/test_pettingzoo_wrapper.py index 58c978d8f10..71b55042c04 100644 --- a/ml-agents-envs/tests/test_pettingzoo_wrapper.py +++ b/ml-agents-envs/tests/test_pettingzoo_wrapper.py @@ -1,3 +1,8 @@ +import numpy as np +import pytest + +from mlagents_envs.base_env import DecisionSteps, TerminalSteps +from mlagents_envs.envs.env_helpers import _unwrap_batch_steps from mlagents_envs.envs.unity_aec_env import UnityAECEnv from mlagents_envs.envs.unity_parallel_env import UnityParallelEnv from simple_test_envs import SimpleEnvironment, MultiAgentEnvironment @@ -6,6 +11,40 @@ NUM_TEST_CYCLES = 100 +def _make_terminal_batch(interrupted): + """A batch with a single terminating agent and no live (decision) agents.""" + obs = [np.array([[1, 2, 3]], dtype=np.float32)] + reward = np.array([1.0]) + interrupted_arr = np.array([interrupted], dtype=bool) + agent_id = np.array([0]) + group_id = np.array([0]) + group_reward = np.array([0.0]) + decision_batch = DecisionSteps([], [], [], [], [], []) + termination_batch = TerminalSteps( + obs, reward, interrupted_arr, agent_id, group_id, group_reward + ) + return decision_batch, termination_batch + + +@pytest.mark.parametrize( + "interrupted,exp_terminated,exp_truncated", + [(True, False, True), (False, True, False)], + ids=["truncated", "terminated"], +) +def test_unwrap_batch_steps_terminated_truncated( + interrupted, exp_terminated, exp_truncated +): + # An interrupted terminal step is a truncation; a non-interrupted one is a + # termination. This is the pettingzoo-side counterpart of the gym split. + (_, _, terminations, truncations, _, _, infos, _) = _unwrap_batch_steps( + _make_terminal_batch(interrupted), "MockBrain" + ) + agent_id = "MockBrain?agent_id=0" + assert terminations[agent_id] is exp_terminated + assert truncations[agent_id] is exp_truncated + assert infos[agent_id]["interrupted"] is exp_truncated + + def test_single_agent_aec(): unity_env = SimpleEnvironment(["test_single"]) env = UnityAECEnv(unity_env) @@ -30,3 +69,18 @@ def test_multi_agent_parallel(): ) env = UnityParallelEnv(unity_env) parallel_api_test(env, num_cycles=NUM_TEST_CYCLES) + + +def test_reset_seed_reseeds_action_spaces(): + # reset(seed=...) on a long-lived wrapper must actually reseed the existing + # action spaces so sampling is reproducible, not just store self._seed. + unity_env = SimpleEnvironment(["test_single"]) + env = UnityAECEnv(unity_env) + + env.reset(seed=1337) + agent = env.possible_agents[0] + first = env.action_space(agent).sample() + env.reset(seed=1337) + second = env.action_space(agent).sample() + + assert np.array_equal(first, second) diff --git a/ml-agents/tests/yamato/scripts/run_gym.py b/ml-agents/tests/yamato/scripts/run_gym.py index 6d698a662b5..f6b54cf2fd9 100644 --- a/ml-agents/tests/yamato/scripts/run_gym.py +++ b/ml-agents/tests/yamato/scripts/run_gym.py @@ -17,7 +17,7 @@ def test_run_environment(env_name): print(str(env)) # Reset the environment - initial_observations = env.reset() + initial_observations, _ = env.reset() if len(env.observation_space.shape) == 1: # Examine the initial vector observation @@ -29,7 +29,8 @@ def test_run_environment(env_name): episode_rewards = 0 while not done: actions = env.action_space.sample() - obs, reward, done, _ = env.step(actions) + obs, reward, terminated, truncated, _ = env.step(actions) + done = terminated or truncated episode_rewards += reward print(f"Total reward this episode: {episode_rewards}") finally: