65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
from torchrl.envs import SerialEnv
|
|
from torchrl.envs.transforms import Compose, StepCounter, Transform, TransformedEnv
|
|
|
|
|
|
def make_env(cfg, transform=None):
|
|
"""
|
|
Note: The returned environment is wrapped in a torchrl.SerialEnv with cfg.rollout_batch_size underlying
|
|
environments. The env therefore returns batches.`
|
|
"""
|
|
|
|
kwargs = {
|
|
"frame_skip": cfg.env.action_repeat,
|
|
"from_pixels": cfg.env.from_pixels,
|
|
"pixels_only": cfg.env.pixels_only,
|
|
"image_size": cfg.env.image_size,
|
|
"num_prev_obs": cfg.n_obs_steps - 1,
|
|
}
|
|
|
|
if cfg.env.name == "simxarm":
|
|
from lerobot.common.envs.simxarm.env import SimxarmEnv
|
|
|
|
kwargs["task"] = cfg.env.task
|
|
clsfunc = SimxarmEnv
|
|
elif cfg.env.name == "pusht":
|
|
from lerobot.common.envs.pusht.env import PushtEnv
|
|
|
|
# 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."
|
|
|
|
clsfunc = PushtEnv
|
|
elif cfg.env.name == "aloha":
|
|
from lerobot.common.envs.aloha.env import AlohaEnv
|
|
|
|
kwargs["task"] = cfg.env.task
|
|
clsfunc = AlohaEnv
|
|
else:
|
|
raise ValueError(cfg.env.name)
|
|
|
|
def _make_env(seed):
|
|
nonlocal kwargs
|
|
kwargs["seed"] = seed
|
|
env = clsfunc(**kwargs)
|
|
|
|
# limit rollout to max_steps
|
|
env = TransformedEnv(env, StepCounter(max_steps=cfg.env.episode_length))
|
|
|
|
if transform is not None:
|
|
# useful to add normalization
|
|
if isinstance(transform, Compose):
|
|
for tf in transform:
|
|
env.append_transform(tf.clone())
|
|
elif isinstance(transform, Transform):
|
|
env.append_transform(transform.clone())
|
|
else:
|
|
raise NotImplementedError()
|
|
|
|
return env
|
|
|
|
return SerialEnv(
|
|
cfg.rollout_batch_size,
|
|
create_env_fn=_make_env,
|
|
create_env_kwargs=[
|
|
{"seed": env_seed} for env_seed in range(cfg.seed, cfg.seed + cfg.rollout_batch_size)
|
|
],
|
|
)
|