Error with env.reset() in robomimic enviroment (in google colab vision-based)

[carloskiyoto]

Hello! First of all thank you for all your work

I am trying to implement a different dataset using the robomimic enviroment (Lift Block)
Since when I'm having troubles with my env.reset to get observations of the enviroment
I used the RobomimicImageWrapper:

#@markdown ### ** Robomimic Environment**
#@markdown Defines `RobomimicImageWrapper`.
#@markdown
from typing import List, Optional
from matplotlib.pyplot import fill
import numpy as np
import gym
from gym import spaces
from omegaconf import OmegaConf
from robomimic.envs.env_robosuite import EnvRobosuite
import sys

class RobomimicImageWrapper(gym.Env):
    def __init__(self,
        env: EnvRobosuite,
        shape_meta: dict,
        init_state: Optional[np.ndarray]=None,
        render_obs_key='agentview_image',
        ):

        self.env = env
        self.render_obs_key = render_obs_key
        self.init_state = init_state
        self.seed_state_map = dict()
        self._seed = None
        self.shape_meta = shape_meta
        self.render_cache = None
        self.has_reset_before = False

        # setup spaces
        action_shape = shape_meta['action']['shape']
        action_space = spaces.Box(
            low=-1,
            high=1,
            shape=action_shape,
            dtype=np.float32
        )
        self.action_space = action_space

        observation_space = spaces.Dict()
        for key, value in shape_meta['obs'].items():
            shape = value['shape']
            min_value, max_value = -1, 1
            if key.endswith('image'):
                min_value, max_value = 0, 1
            elif key.endswith('quat'):
                min_value, max_value = -1, 1
            elif key.endswith('qpos'):
                min_value, max_value = -1, 1
            elif key.endswith('pos'):
                # better range?
                min_value, max_value = -1, 1
            else:
                raise RuntimeError(f"Unsupported type {key}")

            this_space = spaces.Box(
                low=min_value,
                high=max_value,
                shape=shape,
                dtype=np.float32
            )
            observation_space[key] = this_space
        self.observation_space = observation_space

    def debug_print(*args, **kwargs):
        print(*args, **kwargs)
        sys.stdout.flush()

    def get_observation(self, raw_obs=None):
        if raw_obs is None:
            raw_obs = self.env.get_observation()

        self.render_cache = raw_obs[self.render_obs_key]

        obs = dict()
        for key in self.observation_space.keys():
            obs[key] = raw_obs[key]
        return obs

    def seed(self, seed=None):
        np.random.seed(seed=seed)
        self._seed = seed

    def reset(self):
        if self.init_state is not None:
            if not self.has_reset_before:
                # the env must be fully reset at least once to ensure correct rendering
                self.has_reset_before = True

            # always reset to the same state
            # to be compatible with gym
            raw_obs = self.env.reset_to({'states': self.init_state})
            
        elif self._seed is not None:
            # reset to a specific seed
            seed = self._seed
            if seed in self.seed_state_map:
                # env.reset is expensive, use cache
                raw_obs = self.env.reset_to({'states': self.seed_state_map[seed]})
            else:
                # robosuite's initializes all use numpy global random state
                np.random.seed(seed=seed)
                raw_obs = self.env.reset()
                if raw_obs is None:
                    print("Warning: self.env.reset() returned None. Returning an empty dictionary for raw_obs.")
                    raw_obs = {}  # or provide a default value instead of an empty dictionary
                state = self.env.get_state()['states']
                self.seed_state_map[seed] = state
            self._seed = None
        else:
            # random reset
            raw_obs = self.env.reset()
            if raw_obs is None:
                print("Warning: self.env.reset() returned None. Returning an empty dictionary for raw_obs.")
                raw_obs = {}  # or provide a default value instead of an empty dictionary
        # return obs
        obs = self.get_observation(raw_obs)
        return obs
    def step(self, action):
        raw_obs, reward, done, info = self.env.step(action)
        obs = self.get_observation(raw_obs)
        return obs, reward, done, info

    def render(self, mode='rgb_array'):
        if self.render_cache is None:
            raise RuntimeError('Must run reset or step before render.')
        img = np.moveaxis(self.render_cache, 0, -1)
        img = (img * 255).astype(np.uint8)
        return img


def test():
    import os
    from omegaconf import OmegaConf
    cfg_path = os.path.expanduser('/content/diffusion_policy/diffusion_policy/config/task/lift_image.yaml')
    cfg = OmegaConf.load(cfg_path)
    shape_meta = cfg['shape_meta']


    import robomimic.utils.file_utils as FileUtils
    import robomimic.utils.env_utils as EnvUtils
    from matplotlib import pyplot as plt

    dataset_path = os.path.expanduser('/content/image_v141.hdf5')
    env_meta = FileUtils.get_env_metadata_from_dataset(
        dataset_path)

    env = EnvUtils.create_env_from_metadata(
        env_meta=env_meta,
        render=False,
        render_offscreen=False,
        use_image_obs=True,
    )
    

    #print("Env created:", env)

    wrapper = RobomimicImageWrapper(
        env=env,
        shape_meta=shape_meta
    )
    
    wrapper.seed(0)
    obs = wrapper.reset()
    img = wrapper.render()
    plt.imshow(img)





    # states = list()
    # for _ in range(2):
    #     wrapper.seed(0)
    #     wrapper.reset()
    #     states.append(wrapper.env.get_state()['states'])
    # assert np.allclose(states[0], states[1])

    # img = wrapper.render()
    # plt.imshow(img)
    # wrapper.seed()
    # states.append(wrapper.env.get_state()['states'])

when loading test() this error appear:

TypeError                                 Traceback (most recent call last)
[<ipython-input-7-807c2068d7ba>](https://localhost:8080/#) in <cell line: 169>()
    167     plt.imshow(img)
    168 
--> 169 test()
    170 
    171 

3 frames
[<ipython-input-7-807c2068d7ba>](https://localhost:8080/#) in test()
    163 
    164     wrapper.seed(0)
--> 165     obs = wrapper.reset()
    166     img = wrapper.render()
    167     plt.imshow(img)

[<ipython-input-7-807c2068d7ba>](https://localhost:8080/#) in reset(self)
    101                 # robosuite's initializes all use numpy global random state
    102                 np.random.seed(seed=seed)
--> 103                 print(self.env.reset())
    104                 raw_obs = self.env.reset()
    105                 if raw_obs is None:

[/content/diffusion_policy/robomimic/robomimic/envs/env_robosuite.py](https://localhost:8080/#) in reset(self)
    136         """
    137         di = self.env.reset()
--> 138         return self.get_observation(di)
    139 
    140     def reset_to(self, state):

[/content/diffusion_policy/robomimic/robomimic/envs/env_robosuite.py](https://localhost:8080/#) in get_observation(self, di)
    214         ret = {}
    215         for k in di:
--> 216             if (k in ObsUtils.OBS_KEYS_TO_MODALITIES) and ObsUtils.key_is_obs_modality(key=k, obs_modality="rgb"):
    217                 # by default images from mujoco are flipped in height
    218                 ret[k] = di[k][::-1]

TypeError: argument of type 'NoneType' is not iterable

The error happens when self.env.reset() is executed
I will be so grateful if you help me with this problem
Thank you

[carloskiyoto]

I believe the problem is some missing dataset keys, but I can't figure out

[kaustubhsridhar]

I've got the same issue. Pinging @cheng-chi
In the meantime, @carloskiyoto , I think https://github.com/irom-princeton/dppo/blob/main/env/gym_utils/wrapper/robomimic_image.py works.

[AlanxChen]

Using this code to create env:
def create_env(env_meta, shape_meta, enable_render=True): modality_mapping = collections.defaultdict(list) for key, attr in shape_meta['obs'].items(): modality_mapping[attr.get('type', 'low_dim')].append(key) ObsUtils.initialize_obs_modality_mapping_from_dict(modality_mapping) env = EnvUtils.create_env_from_metadata( env_meta=env_meta, render=False, render_offscreen=enable_render, use_image_obs=enable_render, ) return env

Instead of this code to creatre env directly
env = EnvUtils.create_env_from_metadata( env_meta=env_meta, render=False, render_offscreen=False, use_image_obs=True, )