araffin
commented
1 year ago Duplicate of https://github.com/DLR-RM/stable-baselines3/issues/314
The fork is here: https://github.com/MetcalfeTom/stable-baselines3-GPU
EDIT: you will probably need a fork of SB3 that keeps every tensor on the GPU (the fork doesn't seem to do it for vec env).
🚀 Feature
Hello Stable Baselines community,
I am currently working on state representation learning in Robotics, using an observation that consists of a latent vector obtained by encoding simulation images using an Autoencoder's encoder. At each timestep, I process these images on my GPU using the pre-trained Autoencoder's encoder and then pass the resulting latent vector to SB3 as an observation. I have to transfer it to a numpy array on my CPU.
However, I am wondering if there is a way to pass this observation as a Pytorch Tensor instead, to avoid the additional step of transferring the image between the GPU and CPU. Unfortunately, passing the image as the observation is not feasible due to the significant memory required for storing the images.
For context, I am using PPO as the Policy Type.
Thank you in advance!
Motivation
No response
Pitch
No response
Alternatives
No response
Additional context
No response
Checklist