Handing mission space in Babyai env

[Chainesh]

🐛 Bug

Hii,

I'm trying to train BabyAI environment for specific missions during training. The observation space of BabyAI looks like this Dict('direction': Discrete(4), 'image': Box(0, 255, (7, 7, 3), uint8), 'mission': MissionSpace(<function BabyAIMissionSpace._gen_mission at 0x7f813d9da940>, None)) and using RGBImgPartialObsWrapper I've converted the observation Space to Dict('direction': Discrete(4), 'image': Box(0, 255, (56, 56, 3), uint8), 'mission': MissionSpace(<function BabyAIMissionSpace._gen_mission at 0x7aa32616bc40>, None)).

Whether I pass DummyVecEnv or not it throws me the same error TypeError: 'NoneType' object is not iterable due to the fact that mission key is None. How can I train this model considering I want to use MultiInputPolicy after converting the mission into vector embedding and concatenating it with image embedding.

Any help on this would be much appreciated. Let me know if something else is needed to solve this. Thanks :)

Code example

import gymnasium as gym 
from minigrid.wrappers import RGBImgObsWrapper, RGBImgPartialObsWrapper
from stable_baselines3.common.vec_env import VecTransposeImage, DummyVecEnv
from minigrid.wrappers import RGBImgPartialObsWrapper
from stable_baselines3.common.monitor import Monitor
from stable_baselines3 import PPO

class CustomEnv(gym.Env):
    def __init__(self, env, mission):
        self.env = env
        self.observation_space = env.observation_space
        self.action_space = env.action_space
        self.mission = mission

    def reset(self, **kwargs):
        obs = self.env.reset(**kwargs)
        while self.env.mission != self.mission:
            obs = self.env.reset(**kwargs)
        return obs

    def step(self, action):
        return self.env.step(action)

    def render(self, *args, **kwargs):
        return self.env.render(*args, **kwargs)

missions  = ["pick up the red ball"]
for mission in missions:
    env  = gym.make("BabyAI-PickupLoc-v0")#, render_mode = "human" )
    env = RGBImgPartialObsWrapper(env)
    env = CustomEnv(env, mission)
    print(env.observation_space)
    env = Monitor(env)
    env = DummyVecEnv([lambda: env])
    env = VecTransposeImage(env)
    print(env.observation_space)
    model =  PPO("MlpPolicy", env, seed = 42, verbose=1)
    model.learn(50000)

Relevant log output / Error message

Dict('direction': Discrete(4), 'image': Box(0, 255, (56, 56, 3), uint8), 'mission': MissionSpace(<function BabyAIMissionSpace._gen_mission at 0x7aa32616bc40>, None))
Traceback (most recent call last):
  File "/home/scl/ranfom.py", line 43, in <module>
    env = DummyVecEnv([lambda: env])
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/scl/anaconda3/lib/python3.11/site-packages/stable_baselines3/common/vec_env/dummy_vec_env.py", line 46, in __init__
    self.buf_obs = OrderedDict([(k, np.zeros((self.num_envs, *tuple(shapes[k])), dtype=dtypes[k])) for k in self.keys])
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/scl/anaconda3/lib/python3.11/site-packages/stable_baselines3/common/vec_env/dummy_vec_env.py", line 46, in <listcomp>
    self.buf_obs = OrderedDict([(k, np.zeros((self.num_envs, *tuple(shapes[k])), dtype=dtypes[k])) for k in self.keys])
                                                              ^^^^^^^^^^^^^^^^
TypeError: 'NoneType' object is not iterable

System Info

No response

Checklist

• [X] I have checked that there is no similar issue in the repo
• [X] I have read the documentation
• [X] I have provided a minimal and working example to reproduce the bug
• [x] I have checked my env using the env checker
• [X] I've used the markdown code blocks for both code and stack traces.

[qgallouedec]

Hey, please provide the system info As far as I remember, the babyai mission is a str right?

[Chainesh]

Here is the system info as requested:

• OS: Linux-6.5.0-28-generic-x86_64-with-glibc2.35 # 29~22.04.1-Ubuntu SMP PREEMPT_DYNAMIC Thu Apr 4 14:39:20 UTC 2
• Python: 3.11.5
• Stable-Baselines3: 2.2.1
• PyTorch: 2.1.0+cu118
• GPU Enabled: True
• Numpy: 1.24.3
• Cloudpickle: 2.2.1
• Gymnasium: 0.29.1
• OpenAI Gym: 0.26.2

Yes you're correct it should be a string based on the babyai level, but image and direction key are not causing the issue for sure. Even on simply Using this code

env  = gym.make("BabyAI-PickupLoc-v0")
model = PPO("MlpPolicy", env, seed = 42, verbose=1)
model.learn(50000)
env.close()

throws the same error.

[qgallouedec]

Text observation is not supported. Technically, to make it work, you can convert the text to a discrete value, for example. That said, depending on what you're trying to do and show, this solution may not be relevant.

[Chainesh]

Can I convert text into a vector embedding and them might be able to use it?

[qgallouedec]

it should work yes

[Chainesh]

Thanks :)