How do I free system RAM when from_pixels=True in SyncDataCollector?

reinforcement-learning
1

I wrote the following reinforcement learning code in PyTorch, TorchRL. This changes the total_frames and frames_per_batch, but the code is basically based on the PPO tutorial: Google colab

In addition, a 96x96x3 image was obtained from the InvertedDoublePendulum in order to obtain visual data of the environment in the training.

base_env = GymEnv("InvertedDoublePendulum-v4", device=device, frame_skip=frame_skip, from_pixels=True, pixels_only=False)
env = TransformedEnv(
    base_env,
    Compose(
        ObservationNorm(in_keys=["observation"]),
        PermuteTransform((-1, -3, -2), in_keys=["pixels"]),
        Resize(96, 96),
        PermuteTransform((-2, -1, -3), in_keys=["pixels"]),
        ToTensorImage(),
        DoubleToFloat(),
        StepCounter(),
    ),
)

This programme eats up memory as soon as it runs and is forced to terminate. And this eats up even 50 GB of RAM. I suspect this is because SyncDataCollector is dealing with pixel data.

image

However, image data is essential for me, so removing the from_pixels option is unthinkable. The total_frames is also necessary because if the total number of frames is reduced, the learning process is not completely finished. Changing the image size will sooner or later eat up RAM.

I need to save or free the RAM somehow. Do you have any ideas?

2

My fist answer would be that you can reshuffle this and have a better perf by immediately resizing

        ToTensorImage(),
        Resize(96, 96),
        ObservationNorm(in_keys=["observation"]),
        DoubleToFloat(),
        StepCounter(),

Otherwise you are normalizing big images (which takes RAM) when you could do that after resize. Does that help?

3

Thanks for the suggestion. I have tried it with Google Colab, but the problem of memory build up remains.

This is a cut-out of a Google Colab screen. The green indicator just to the left of for i , tensordict... indicates that this part is running.

image
image1588×265 21.8 KB

At the time of collecting data with this SyncDataCollector (for i, tensordict_data...) and at the env.transform.init_stats, memory usage increases and is accumulated as it is. As a result, the system RAM is wasted, as shown on the right, and the programme eventually stops.

Not sure if this is relevant, here is the version of each package, for example. It is simply a pip install run on the Google Colab:

import torchrl, numpy, sys
print(torchrl.__version__, numpy.__version__, sys.version, sys.platform)

0.5.0 1.26.4 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0] linux
4

Can I ask why are you importing split_trajectories within the loop?
It seems that loads the memory at every iteration. I don’t see any increase in RAM if I take it away.
LMK if that makes sense

5

That is the module I was going to use in an interations. I tried to remove that, but the memory still keeps growing.

image
image1587×295 20.6 KB