12345678910111213141516171819202122232425262728293031 |
- from gymnasium.envs.classic_control.cartpole import CartPoleEnv
- from gymnasium.utils import EzPickle
- import numpy as np
- from ray.rllib.env.apis.task_settable_env import TaskSettableEnv
- class CartPoleMassEnv(CartPoleEnv, EzPickle, TaskSettableEnv):
- """CartPoleMassEnv varies the weights of the cart and the pole."""
- def sample_tasks(self, n_tasks):
- # Sample new cart- and pole masses (random floats between 0.5 and 2.0
- # (cart) and between 0.05 and 0.2 (pole)).
- cart_masses = np.random.uniform(low=0.5, high=2.0, size=(n_tasks, 1))
- pole_masses = np.random.uniform(low=0.05, high=0.2, size=(n_tasks, 1))
- return np.concatenate([cart_masses, pole_masses], axis=-1)
- def set_task(self, task):
- """
- Args:
- task (Tuple[float]): Masses of the cart and the pole.
- """
- self.masscart = task[0]
- self.masspole = task[1]
- def get_task(self):
- """
- Returns:
- Tuple[float]: The current mass of the cart- and pole.
- """
- return np.array([self.masscart, self.masspole])
|