Environment examples

Running the Gymnasium environments that wrap Sim for reinforcement learning. See Gymnasium Environments for the available environments.

---

Figure-eight environment

Run the figure-eight Gymnasium environment with a constant upward attitude command.

The env wraps Sim behind the standard Gymnasium VectorEnv interface.

examples/environments/figure8.py

"""Run the figure-eight Gymnasium environment with a constant upward attitude command.

The env wraps `Sim` behind the standard Gymnasium `VectorEnv` interface.
"""

import gymnasium
import jax.numpy as jnp
import numpy as np
from gymnasium.wrappers.vector import JaxToNumpy  # , JaxToTorch

from crazyflow.envs import NormalizeActions  # noqa: F401
from crazyflow.utils import enable_cache


def main():
    enable_cache()
    # Create environment that contains a figure eight trajectory. You can parametrize the
    # observation space, i.e., which part of the trajectory is contained in the observation. Please
    # refer to the documentation of the environment for more information.
    envs = gymnasium.make_vec(
        "DroneFigureEightTrajectory-v0",
        num_envs=20,
        freq=50,
        n_samples=10,
        samples_dt=0.1,
        trajectory_time=10.0,
    )

    # NormalizeActions wrapper to clip the actions to [-1, 1] and rescale them for use with common
    # DRL libraries.
    envs = NormalizeActions(envs)
    envs = JaxToNumpy(envs)

    # dummy action for going up (in attitude control)
    action = np.zeros((20, 4), dtype=np.float32)
    action[..., 3] = 0.3

    obs, info = envs.reset()
    # Step through the environment
    for _ in range(1_000):
        # Prevent alignment warnings. Related issue: https://github.com/jax-ml/jax/issues/29810
        # TODO: Remove once https://github.com/jax-ml/jax/pull/29963 is merged.
        action = np.asarray(jnp.asarray(action))
        observation, reward, terminated, truncated, info = envs.step(action)
        envs.render()

    envs.close()


if __name__ == "__main__":
    main()

---

Reach-position environment

Run the reach-position Gymnasium environment with a constant action.

The JaxToNumpy wrapper converts between NumPy actions and JAX observations.

examples/environments/gymnasium_env.py

"""Run the reach-position Gymnasium environment with a constant action.

The `JaxToNumpy` wrapper converts between NumPy actions and JAX observations.
"""

import gymnasium
import jax.numpy as jnp
import numpy as np
from gymnasium.wrappers.vector import JaxToNumpy  # , JaxToTorch

import crazyflow  # noqa: F401, register gymnasium envs
from crazyflow.utils import enable_cache


def main():
    enable_cache()
    envs = gymnasium.make_vec("DroneReachPos-v0", num_envs=20, freq=50, max_episode_time=2)

    # This wrapper converts numpy actions to jax arrays and jax observations to numpy arrays. To
    # convert to torch, use JaxToTorch.
    envs = JaxToNumpy(envs)

    # Dummy action for going up (in attitude control)
    action = np.zeros((20, 4), dtype=np.float32)
    action[..., 0] = 0.4

    # Environments provide reset parameters that can be used to set the initial state of the
    # environment.
    obs, info = envs.reset(
        options={
            "goal_pos_min": np.array([-1.0, 1.0, 1.0]),
            "goal_pos_max": np.array([-1.0, 1.0, 1.0]),
        }
    )

    # Step through the environment
    for _ in range(1_000):
        # Prevent alignment warnings. Related issue: https://github.com/jax-ml/jax/issues/29810
        # TODO: Remove once https://github.com/jax-ml/jax/pull/29963 is merged.
        action = np.asarray(jnp.asarray(action))
        observation, reward, terminated, truncated, info = envs.step(action)
        envs.render()
    envs.close()


if __name__ == "__main__":
    main()