58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
from torchrl.envs.transforms import StepCounter, TransformedEnv
|
|
|
|
from lerobot.common.envs.pusht import PushtEnv
|
|
from lerobot.common.envs.simxarm import SimxarmEnv
|
|
from lerobot.common.envs.transforms import Prod
|
|
|
|
|
|
def make_env(cfg):
|
|
kwargs = {
|
|
"frame_skip": cfg.action_repeat,
|
|
"from_pixels": cfg.from_pixels,
|
|
"pixels_only": cfg.pixels_only,
|
|
"image_size": cfg.image_size,
|
|
}
|
|
|
|
if cfg.env == "simxarm":
|
|
kwargs["task"] = cfg.task
|
|
clsfunc = SimxarmEnv
|
|
elif cfg.env == "pusht":
|
|
clsfunc = PushtEnv
|
|
else:
|
|
raise ValueError(cfg.env)
|
|
|
|
env = clsfunc(**kwargs)
|
|
|
|
# limit rollout to max_steps
|
|
env = TransformedEnv(env, StepCounter(max_steps=cfg.episode_length))
|
|
|
|
if cfg.env == "pusht":
|
|
# to ensure pusht is in [0,255] like simxarm
|
|
env.append_transform(Prod(in_keys=[("observation", "image")], prod=255.0))
|
|
|
|
return env
|
|
|
|
|
|
# def make_env(env_name, frame_skip, device, is_test=False):
|
|
# env = GymEnv(
|
|
# env_name,
|
|
# frame_skip=frame_skip,
|
|
# from_pixels=True,
|
|
# pixels_only=False,
|
|
# device=device,
|
|
# )
|
|
# env = TransformedEnv(env)
|
|
# env.append_transform(NoopResetEnv(noops=30, random=True))
|
|
# if not is_test:
|
|
# env.append_transform(EndOfLifeTransform())
|
|
# env.append_transform(RewardClipping(-1, 1))
|
|
# env.append_transform(ToTensorImage())
|
|
# env.append_transform(GrayScale())
|
|
# env.append_transform(Resize(84, 84))
|
|
# env.append_transform(CatFrames(N=4, dim=-3))
|
|
# env.append_transform(RewardSum())
|
|
# env.append_transform(StepCounter(max_steps=4500))
|
|
# env.append_transform(DoubleToFloat())
|
|
# env.append_transform(VecNorm(in_keys=["pixels"]))
|
|
# return env
|