34 lines
889 B
Python
34 lines
889 B
Python
from mp_env_api.mp_wrappers.dmp_wrapper import DmpWrapper
|
|
from mp_env_api.mp_wrappers.detpmp_wrapper import DetPMPWrapper
|
|
import gym
|
|
from gym.vector.utils import write_to_shared_memory
|
|
import sys
|
|
|
|
|
|
def make_env(env_id, seed, rank):
|
|
env = gym.make(env_id)
|
|
env.seed(seed + rank)
|
|
return lambda: env
|
|
|
|
|
|
def make_contextual_env(env_id, context, seed, rank):
|
|
env = gym.make(env_id, context=context)
|
|
env.seed(seed + rank)
|
|
return lambda: env
|
|
|
|
|
|
def make_dmp_env(**kwargs):
|
|
name = kwargs.pop("name")
|
|
_env = gym.make(name)
|
|
for wrapper in kwargs.pop("wrappers"):
|
|
_env = wrapper(_env)
|
|
return DmpWrapper(_env, **kwargs.get("mp_kwargs"))
|
|
|
|
|
|
def make_detpmp_env(**kwargs):
|
|
name = kwargs.pop("name")
|
|
_env = gym.make(name)
|
|
for wrapper in kwargs.pop("wrappers"):
|
|
_env = wrapper(_env)
|
|
return DetPMPWrapper(_env, **kwargs.get("mp_kwargs"))
|