# Training on multiple environment

**URL:** <https://discuss.ray.io/t/training-on-multiple-environment/9289>\
**Category:** Offline RL\
**Created:** [February 7, 2023, 9:42pm UTC](https://discuss.ray.io/t/training-on-multiple-environment/9289 "2023-02-07T21:42:26Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![Erfan\_Asaadi](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/erfan_asaadi/32/3855_2.png) [@Erfan\_Asaadi](https://discuss.ray.io/u/Erfan_Asaadi)\
**Post date:** [February 7, 2023, 9:42pm UTC](https://discuss.ray.io/t/training-on-multiple-environment/9289/1 "2023-02-07T21:42:26Z")

</div>

Hi,  
I was wondering if anybody has a suggestion/comment on the following problem:  
I am trying to train an agent in multiple environment at the same time. In the other word, at the end of the episode I a need to switch between the environments.  
Except some environmental parameters, the environments are similar, so I am going to switch between the environmental parameters when the agent “reset” the environments. It is fine when there is no parallel computing, but now that I am using rllib, I do not know if it cause an issue, e.g. one process looks at the old environment when another process looks at the new environment.  
Any suggestions is highly appreciated

---

<div class="post-metadata">

**Author:** ![sven1977](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/sven1977/32/53_2.png) [@sven1977](https://discuss.ray.io/u/sven1977)\
**Post date:** [February 9, 2023, 9:17am UTC](https://discuss.ray.io/t/training-on-multiple-environment/9289/2 "2023-02-09T09:17:06Z")

</div>

Hey @Erfan_Asaadi , great question.  
You could use one of our callbacks for this task. However, for this to work, you should have your environment class designed such that it can take all the different tasks you would like to simulate within a single gym.Env class (basically have a switch\_task() method inside your env):

```auto
from ray.rllib.algorithms.callbacks import DefaultCallbacks
from ray.rllib.algorithms.ppo import PPOConfig

class MyCallbacks(DefaultCallbacks):
    def on_episode_created(
        self,
        *,
        worker: "RolloutWorker",
        base_env: BaseEnv,
        policies: Dict[PolicyID, Policy],
        env_index: int,
        episode: Union[Episode, EpisodeV2],
        **kwargs
    ):
        # get the respective sub-environment that is about to be reset.
        sub_environment_to_change = worker.base_env.get_sub_environments()[env_index]
        # switch its task
        sub_environment_to_change.switch_task(task=...)

# Then, in your AlgorithmConfig object:
config = PPOConfig().callbacks(MyCallbacks)

```

Make sure your gym.Env subclass implements the `switch_task` method taking whatever arguments needed to switch between tasks.

---

<div class="post-metadata">

**Author:** ![Erfan\_Asaadi](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/erfan_asaadi/32/3855_2.png) [@Erfan\_Asaadi](https://discuss.ray.io/u/Erfan_Asaadi)\
**Post date:** [February 14, 2023, 1:05am UTC](https://discuss.ray.io/t/training-on-multiple-environment/9289/3 "2023-02-14T01:05:58Z")

</div>

import gym  
from gym import spaces  
import numpy as np  
import test\_exercises  
import ray  
from ray.rllib.agents.ppo import PPOTrainer, DEFAULT\_CONFIG

ray.init(ignore\_reinit\_error=True, log\_to\_driver=False)

class ChainEnv(gym.Env):  
def **init** (self, env\_config = None):  
env\_config = env\_config or {}  
print(env\_config)  
self.n = env\_config.get(“n”, 6)  
self.reward\_ratio = .1  
self.small\_reward = env\_config.get(“small”, 2 \* self.reward\_ratio) # payout for ‘backwards’ action  
self.large\_reward = env\_config.get(“large”, 10 \* self.reward\_ratio) # payout at end of chain for ‘forwards’ action  
self.\_horizon = self.n  
self.\_counter = 0 # For terminating the episode  
self.\_setup\_spaces()  
self.cnt = 0  
self.task\_num = 0

```
def _setup_spaces(self):
    self.action_space = spaces.Discrete(2)
    self.observation_space = spaces.Discrete(6)

def step(self, action):
    assert self.action_space.contains(action)
    if action == 1: # 'backwards': go back to the beginning, get small reward
        reward = self.small_reward
        self.state = 0
    elif self.state < self.n - 1: # 'forwards': go up along the chain
        reward = 0
        self.state += 1
    else: # 'forwards': stay at the end of the chain, collect large reward
        reward = self.large_reward
    self._counter += 1
    done = self._counter >= self._horizon
    return self.state, reward, done, {}

def reset(self):
    self.cnt += 1
    self.state = 0
    self._counter = 0
    return self.state    

def switch_task(self):       
    if self.task_num % 2 == 0:
      self.reward_ratio = .1
    else:
      self.reward_ratio = 1
    self.task_num += 1
    self.small_reward = env_config.get("small", 2 * self.reward_ratio) # payout for 'backwards' action
    self.large_reward = env_config.get("large", 10 * self.reward_ratio) # payout at end of chain for 'forwards' action 

```

from ray.rllib.algorithms.dqn import DQNConfig  
from ray.rllib.algorithms.callbacks import DefaultCallbacks  
from ray.rllib.evaluation.episode import Episode

class MyCallbacks(DefaultCallbacks):  
def on\_episode\_created(  
self, \*, worker, base\_env, policies, env\_index, episode, \*\*kwargs):  
# get the respective sub-environment that is about to be reset.  
sub\_environment\_to\_change = worker.base\_env.get\_sub\_environments()[env\_index]  
# switch its task  
sub\_environment\_to\_change.switch\_task()

config = DQNConfig()  
config = config.rollouts(num\_rollout\_workers=4).callbacks(MyCallbacks)  
simple\_trainer = config.environment(env=ChainEnv).build()  
for \_ in range(1):  
A = simple\_trainer.train()