Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

How to use image on cuda with gym wrapper #742

Open
chrisgao99 opened this issue Jul 16, 2024 · 0 comments
Open

How to use image on cuda with gym wrapper #742

chrisgao99 opened this issue Jul 16, 2024 · 0 comments

Comments

@chrisgao99
Copy link

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?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant