How to use image on cuda with gym wrapper

[chrisgao99]

Hello,

I'm trying to train image input rl policy with stable baseline3 and I have this wrapper to make image into right shape. I tried to keep the image stay on gpu to make the training faster by set image_on_cuda to be True but the gym box doesn't accept cupy array or torch tensor.

class MyWrapper(gym.Wrapper):
    def __init__(self, env, seed):
        super().__init__(env)
        self.observation_space = Dict({
            "image": Box(-0.0, 1.0, (3,224,224), dtype=np.float32),
            "state": self.env.observation_space['state'] #BOX -1,1
        })
        self.env_seed = seed
        
    def reset(self, **kwargs):
        kwargs.pop('options', None)        
        kwargs['seed'] = np.random.randint(0,10)
        obs,info = self.env.reset(**kwargs)
        print(type(obs['image']))
        obs['image'] = torch.tensor(obs['image'],device='cuda',dtype=torch.float32)
        obs['image'] = obs['image'].squeeze(-1).permute(2, 0, 1)
        obs['image'] = obs['image'].cpu().numpy()   
        return obs,info
    
    def step(self, action):
        obs, reward, terminate,truncated, info = self.env.step(action)
        # obs['image'] = obs['image'].get()
        obs['image'] = torch.tensor(obs['image'],device='cuda',dtype=torch.float32)
        obs['image'] = obs['image'].squeeze(-1).permute(2, 0, 1)
        obs['image'] = obs['image'].cpu().numpy()
        return obs, reward, terminate, truncated, info

My question is is there any other way to make image on cuda work without transferring the data back to cpu. If not, would it be better that I turn off image_on_cuda and just use cpu to do the interaction with environment?