supporting setting _N_GOALS

reginald-mclean · Nov 26, 2024 · cc018e1 · cc018e1
1 parent 2610866
commit cc018e1
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/metaworld/__init__.py b/metaworld/__init__.py
@@ -550,6 +550,10 @@ def _mt_bench_vector_entry_point(
         *args,
         **lamb_kwargs,
     ):
+        if 'num_goals' in lamb_kwargs:
+            global _N_GOALS
+            _N_GOALS = lamb_kwargs['num_goals']
+            del lamb_kwargs['num_goals']
         return make_mt_envs(  # type: ignore
             mt_bench,
             seed=seed,