49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
import gymnasium as gym
|
|
|
|
|
|
def make_env(cfg, num_parallel_envs=0) -> gym.Env | gym.vector.SyncVectorEnv:
|
|
"""
|
|
Note: When `num_parallel_envs > 0`, this function returns a `SyncVectorEnv` which takes batched action as input and
|
|
returns batched observation, reward, terminated, truncated of `num_parallel_envs` items.
|
|
"""
|
|
kwargs = {}
|
|
|
|
if cfg.env.name == "simxarm":
|
|
import gym_xarm # noqa: F401
|
|
|
|
assert cfg.env.task == "lift"
|
|
env_fn = lambda: gym.make(
|
|
"gym_xarm/XarmLift-v0",
|
|
render_mode="rgb_array",
|
|
max_episode_steps=cfg.env.episode_length,
|
|
**kwargs,
|
|
)
|
|
elif cfg.env.name == "pusht":
|
|
import gym_pusht # noqa: F401
|
|
|
|
# assert kwargs["seed"] > 200, "Seed 0-200 are used for the demonstration dataset, so we don't want to seed the eval env with this range."
|
|
kwargs.update(
|
|
{
|
|
"obs_type": "pixels_agent_pos",
|
|
"render_action": False,
|
|
}
|
|
)
|
|
env_fn = lambda: gym.make( # noqa: E731
|
|
"gym_pusht/PushTPixels-v0",
|
|
render_mode="rgb_array",
|
|
max_episode_steps=cfg.env.episode_length,
|
|
**kwargs,
|
|
)
|
|
elif cfg.env.name == "aloha":
|
|
kwargs["task"] = cfg.env.task
|
|
else:
|
|
raise ValueError(cfg.env.name)
|
|
|
|
if num_parallel_envs == 0:
|
|
# non-batched version of the env that returns an observation of shape (c)
|
|
env = env_fn()
|
|
else:
|
|
# batched version of the env that returns an observation of shape (b, c)
|
|
env = gym.vector.SyncVectorEnv([env_fn for _ in range(num_parallel_envs)])
|
|
return env
|