Repository: Farama-Foundation/Metaworld
Issue: Vectorizing Envs over Many Workers Results in Memory Overflow
Root cause: Algorithmic Error

Currently, I'm using RLlib for running metaworld envs, where each worker runs many vectorized instances of the environment. 

Tried running MAML/ProMP with 40 workers and 20 envs/worker on one of the Metaenvs (Push).  Rllib can train on this for couple iterations before crashing due to memory overflow (can't allocate more memory). Not sure on what exactly is the issue, but do you have some leads on what could be the issue? I was thinking that it might be a memory leak, but trying this on a lower # of workers resulted in worse training overall but most importantly no crashing.

## tl;dr
The issue is that 3D simulation has a fairly large memory footprint, and most currently-available implementations of RL algorithms aren't designed to solve challenges as large as metaworld.

## It's a feature, not a bug
We're aware of this issue and I personally consider it a feature, not a bug.

A benchmark this large challenges researchers to develop algorithms which can run in a feasible memory and time footprint while solving realistic challenges. metaworld is right at the edge of what's possible.

## How we did it
The MAML/ProMP evaluations you will find in our forthcoming arXiv **paper require approximately 1.5TB of memory to run**, assuming you are using [this implementation](https://github.com/jonasrothfuss/ProMP/tree/full_code) and standard hyperparameters for sampling. For ML45/MT50 we use 45/50 workers (i.e. meta-batch size equal to number of environments) and 10 environments per worker. We ran them using Google Compute Engine instance type `m1-ultramem-80`.

Parallel sampling trades space for time, and 1.5TB reflects a 10x parallelization of the environments. If we maintain the invariant the algorithm must see each environment at least once per meta-batch, then **the minimum memory footprint for running ML45/MT50 is around 150GB (3GB/environment)**, and this is what we observed using implementations which don't parallelize environments.

## Is it possible to make it smaller?
I'm sure there are optimizations we could apply to our MuJoCo models which might bring the per-environment footprint down by some integer multiple, perhaps 2 or 3x. If you come across any, please send them as a PR! I doubt there are straightforward optimization which would cut the per-env footprint by an order of magnitude.

The only optimization likely to drive the MuJoCo memory footprint much lower than 150GB would be to merge all of the environments into a single "World" (MuJoCo environment) and then dynamically reconfigure the world during sampling to represent the needed environment.

We didn't do this (yet, and maybe ever) for two reasons:
1. The `gym.Env` interface which most of today's algorithms rely on has no API for mutating an environment.
    To define our benchmark in this way, we would have to create such an API, and then algorithm authors would have to change their algorithms to implement that API in order to sample from the benchmark. We felt this asks too much up-front of most authors, who simply want to run an existing algorithm on a new benchmark. We could provide a wrapper to do this for them, but essentially that's writing a meta-sampler, and sampling is an important detail of algorithm implementation.
1. Most algorithms today prefer to sample from every environment (or a large subset thereof) in every meta-batch.
    This implies that even with a mutable `World` API, most algorithms would keep 45 or 50 copies of the `World` resident in memory. Sure, they could keep only 1 and mutate it 50 times per meta-batch, but that would mean their training process incurs a 50x slowdown compared to one-worker-per-env parallelized sampling. This seems even less feasible than needing 150GB of memory to run MT45/MT50, because state-of-the-art algorithms already require several days to converge using the parallelized version.

I'm marking this as closed. Please feel free to re-open if you have any new questions, or find an optimization.

@ryanjulian It seems like Michael is running to this issue with a single one of the Metaenvs, while you are referring to the setting where you run MT50/ML45. The memory footprint shouldn't be large for a single env.

Separately, for ML45/MT50, it should be possible to get a lower memory footprint by subsampling environments.

@ryanjulian @cbfinn Thanks for the quick reply! Yes, just running ProMP on one of the Metaenvs, where the tasks are the goals, so it is not MT50/ML45. The ProMP implementation I'm using is similar to Jonas's implementation but ported onto Rllib on a fork. Just wondering how much memory does it take to run a single env on your machines?



I think the ProMP implementation MAML, instead of meta-world, has a memory leak, which I could not locate before. @ryanjulian @michaelzhiluo @cbfinn 

ah, my apologies for the misunderstanding. reopening.

my machine takes about 3GB to run a single env. 

@zhanpenghe Not using Jonas's ProMP, but a Rllib implementation of ProMP on a private fork. Pretty sure there is no memory leak on Rllib since it works on many other envs, from Mujoco, Atari, Vizdoom, etc. and Ray has not received a git issue about this in quite a while.

@ryanjulian On your machine, how much memory would HalfCheetah or some RandDir/Velocity variant of it take?


@michaelzhiluo I'm not sure, but I think @deq2 or @tianheyu927 might have a data point for you.

I'm not sure about HalfCheetah, but running a single env in metaworld shouldn't be memory-intensive.

Ryan mentioned that a single MetaWorld env is around 3 GB, which, when vectorized and scaled to multiple CPUs, can heavily consume machine memory. Is it expected that 3 GB is the average env size when compared to Mujoco envs?

I ran metaworld with garage on MT50 before without any memory leak issue. Also the usage for simple benchmark like ML1 or single task should not have high memory usage.

I'm seeing peak memory usage of around 1.1GB during our test suite.

Assuming your workers also contain some memory overhead (a bunch of Python packages and torch or TF), I think 1.5GB/env is a fair expectation.

I'm not sure what the "average" MuJoCo environment footprint is. I ran HalfCheetah on the latest copy of garage and got a minimum footprint of 500-750MB. IME MuJoCo runtime and memory footprint to scales with number of collision bodies, and the metaworld environments have quite a few of those.

@ryanjulian  @tianheyu927  @zhanpenghe  Had some time to run some simple experiments to check memory usage. Experiment code is a below:
```
from metaworld.benchmarks import ML1 
envs = []
for i in range(0, 400):
      envs.append(ML1.get_train_tasks('pick-place-v1'))
      envs[i].reset()
```
I used htop while running this to track memory usage over time as the environments instantiated.
Calling `envs[i].reset()` for the first time takes around 450-500 MB per meta-world environment. Also tested this on Mujoco HalfCheetah-v2--this is around 4.2MB per environment. Tried this on FetchPush-v1 for the Gym environments--this took around 7.7MB per environment.

Do you know what is causing so much memory to be taken per environment? It was mentioned earlier that there are a lot more contact bodies, but I am confused to whether this would contribute to 100x in memory. Also, what did you run to get your 500 MB results for Mujoco envs and 3 GB results for minimum footprint (many environments/run)?

@michaelzhiluo I honestly don't know, and might need to run a memory profiler or ask the folks at MuJoCo to find out.

Do you mind trying the same test with the most complex models in gym, which I suspect are the shadow hand models (e.g. `HandManipulateBlock-v0`)?

@michaelzhiluo I was able to run some tests locally and reproduce your results -- each env takes 500MB of resident memory when even complex MuJoCo models from gym take up ~10MB.

I'm not sure what the cause is here. I tried turning off all of the Sawyer meshes (on the theory that they are very detailed), but got exactly the same behavior.


@ryanjulian Good to know that this bug persists on different machines! For HandManipulate environment, I averaged around 11 MB per env, so the difference between other gym envs is not huge.

What are all possible causes that could contribute to high mem/env?

Quick update:
* I loaded some metaworld modules using the `simulate` tool which ships with MuJoCo, and got memory usage in the 350-400MB range. This is compared to complex models (e.g. `humanoid100.xml` and `particle.xml`) which take up 10-40MB.
* metaworld models are pretty small in the metrics the MuJoCo profiler tracks, such as # of bodies, # of constraints, and # of state axes, so I don't think it's an issue of intrinsic complexity
* A preliminary test of removing textures and meshes didn't make any difference in memory, but I want to repeat that one since I'm not sure I caught every mesh. 

At this point I think the culprits might be (in order of likelihood):
1. Use of high-quality textures and meshes
1. Solver/engine configuration which subtly balloons memory usage
1. Degenerate model configuration, such as a loop in includes

I can eliminate (1) and (3) shortly by trying some tests with the model files. If the culprit is (2) then I think we will have to reach out to the MuJoCo authors for help.

Hi @ryanjulian , thanks for your hardwork! Did you find that memory consumption decreased significantly via fixing (1) and (3)?

To be clear, the model configurations were never degenerate (3) so there's nothing to fix. I included it in my list as something to check.

(1) contributes to memory usage but only at a fixed overhead. We don't have 500MB of textures and meshes

After consulting with some folks who really know MuJoCo, it turns out the problem is likely (2). I will test a fix over the next couple days. If you want to jump ahead, you need to reduce the parameter `njmax` in the mujoco XML models.

@michaelzhiluo see #14 for a fix.

I haven't tested this for a long-running training job, so it's possible MuJoCo will estimate memory needs incorrectly and fail mid-training. In that case we will have to pick some sane but conservative static allocation.

Please give it a try and see if it works for you now.

Just tested this out with Pick-Place-v1; around 48 MB per environment. A lot better than before; thank you for your hardwork :)