123456789101112131415161718192021222324252627282930313233 |
- from gymnasium.envs.classic_control.pendulum import PendulumEnv
- from gymnasium.utils import EzPickle
- import numpy as np
- from ray.rllib.env.apis.task_settable_env import TaskSettableEnv
- class PendulumMassEnv(PendulumEnv, EzPickle, TaskSettableEnv):
- """PendulumMassEnv varies the weight of the pendulum
- Tasks are defined to be weight uniformly sampled between [0.5,2]
- """
- def sample_tasks(self, n_tasks):
- # Sample new pendulum masses (random floats between 0.5 and 2).
- return np.random.uniform(low=0.5, high=2.0, size=(n_tasks,))
- def set_task(self, task):
- """
- Args:
- task: Task of the meta-learning environment (here: mass of
- the pendulum).
- """
- # self.m is the mass property of the pendulum.
- self.m = task
- def get_task(self):
- """
- Returns:
- float: The current mass of the pendulum (self.m in the PendulumEnv
- object).
- """
- return self.m
|