tangger/lerobot_piper
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

test_envs are passing

Browse Source
This commit is contained in:
Cadene
2024-04-05 23:27:12 +00:00
parent 5eff40b3d6
commit 44656d2706
7 changed files with 91 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
lerobot/common/datasets/pusht.py
View File

@@ -6,9 +6,9 @@ import pygame
import pymunk
import torch
import tqdm
from gym_pusht.envs.pusht import pymunk_to_shapely
from lerobot.common.datasets.utils import download_and_extract_zip, load_data_with_delta_timestamps
from lerobot.common.envs.pusht.pusht_env import pymunk_to_shapely
from lerobot.common.policies.diffusion.replay_buffer import ReplayBuffer as DiffusionPolicyReplayBuffer
# as define in env

6
lerobot/common/envs/aloha/__init__.py
View File

@@ -4,6 +4,9 @@ register(
id="gym_aloha/AlohaInsertion-v0",
entry_point="lerobot.common.envs.aloha.env:AlohaEnv",
max_episode_steps=300,
# Even after seeding, the rendered observations are slightly different,
# so we set `nondeterministic=True` to pass `check_env` tests
nondeterministic=True,
kwargs={"obs_type": "state", "task": "insertion"},
)
@@ -11,5 +14,8 @@ register(
id="gym_aloha/AlohaTransferCube-v0",
entry_point="lerobot.common.envs.aloha.env:AlohaEnv",
max_episode_steps=300,
# Even after seeding, the rendered observations are slightly different,
# so we set `nondeterministic=True` to pass `check_env` tests
nondeterministic=True,
kwargs={"obs_type": "state", "task": "transfer_cube"},
)

38
lerobot/common/envs/aloha/env.py
View File

@@ -16,7 +16,6 @@ from lerobot.common.envs.aloha.tasks.sim_end_effector import (
TransferCubeEndEffectorTask,
)
from lerobot.common.envs.aloha.utils import sample_box_pose, sample_insertion_pose
from lerobot.common.utils import set_global_seed
class AlohaEnv(gym.Env):
@@ -55,15 +54,20 @@ class AlohaEnv(gym.Env):
elif self.obs_type == "pixels_agent_pos":
self.observation_space = spaces.Dict(
{
"pixels": spaces.Box(
low=0,
high=255,
shape=(self.observation_height, self.observation_width, 3),
dtype=np.uint8,
"pixels": spaces.Dict(
{
"top": spaces.Box(
low=0,
high=255,
shape=(self.observation_height, self.observation_width, 3),
dtype=np.uint8,
)
}
),
"agent_pos": spaces.Box(
low=np.array([-1] * len(JOINTS)), # ???
high=np.array([1] * len(JOINTS)), # ???
low=-np.inf,
high=np.inf,
shape=(len(JOINTS),),
dtype=np.float64,
),
}
@@ -89,21 +93,21 @@ class AlohaEnv(gym.Env):
if "transfer_cube" in task_name:
xml_path = ASSETS_DIR / "bimanual_viperx_transfer_cube.xml"
physics = mujoco.Physics.from_xml_path(str(xml_path))
task = TransferCubeTask(random=False)
task = TransferCubeTask()
elif "insertion" in task_name:
xml_path = ASSETS_DIR / "bimanual_viperx_insertion.xml"
physics = mujoco.Physics.from_xml_path(str(xml_path))
task = InsertionTask(random=False)
task = InsertionTask()
elif "end_effector_transfer_cube" in task_name:
raise NotImplementedError()
xml_path = ASSETS_DIR / "bimanual_viperx_end_effector_transfer_cube.xml"
physics = mujoco.Physics.from_xml_path(str(xml_path))
task = TransferCubeEndEffectorTask(random=False)
task = TransferCubeEndEffectorTask()
elif "end_effector_insertion" in task_name:
raise NotImplementedError()
xml_path = ASSETS_DIR / "bimanual_viperx_end_effector_insertion.xml"
physics = mujoco.Physics.from_xml_path(str(xml_path))
task = InsertionEndEffectorTask(random=False)
task = InsertionEndEffectorTask()
else:
raise NotImplementedError(task_name)
@@ -116,10 +120,10 @@ class AlohaEnv(gym.Env):
if self.obs_type == "state":
raise NotImplementedError()
elif self.obs_type == "pixels":
obs = raw_obs["images"]["top"].copy()
obs = {"top": raw_obs["images"]["top"].copy()}
elif self.obs_type == "pixels_agent_pos":
obs = {
"pixels": raw_obs["images"]["top"].copy(),
"pixels": {"top": raw_obs["images"]["top"].copy()},
"agent_pos": raw_obs["qpos"],
}
return obs
@@ -129,14 +133,14 @@ class AlohaEnv(gym.Env):
# TODO(rcadene): how to seed the env?
if seed is not None:
set_global_seed(seed)
self._env.task.random.seed(seed)
self._env.task._random = np.random.RandomState(seed)
# TODO(rcadene): do not use global variable for this
if "transfer_cube" in self.task:
BOX_POSE[0] = sample_box_pose() # used in sim reset
BOX_POSE[0] = sample_box_pose(seed) # used in sim reset
elif "insertion" in self.task:
BOX_POSE[0] = np.concatenate(sample_insertion_pose()) # used in sim reset
BOX_POSE[0] = np.concatenate(sample_insertion_pose(seed)) # used in sim reset
else:
raise ValueError(self.task)

14
lerobot/common/envs/aloha/utils.py
View File

@@ -1,26 +1,30 @@
import numpy as np
def sample_box_pose():
def sample_box_pose(seed=None):
x_range = [0.0, 0.2]
y_range = [0.4, 0.6]
z_range = [0.05, 0.05]
rng = np.random.RandomState(seed)
ranges = np.vstack([x_range, y_range, z_range])
cube_position = np.random.uniform(ranges[:, 0], ranges[:, 1])
cube_position = rng.uniform(ranges[:, 0], ranges[:, 1])
cube_quat = np.array([1, 0, 0, 0])
return np.concatenate([cube_position, cube_quat])
def sample_insertion_pose():
def sample_insertion_pose(seed=None):
# Peg
x_range = [0.1, 0.2]
y_range = [0.4, 0.6]
z_range = [0.05, 0.05]
rng = np.random.RandomState(seed)
ranges = np.vstack([x_range, y_range, z_range])
peg_position = np.random.uniform(ranges[:, 0], ranges[:, 1])
peg_position = rng.uniform(ranges[:, 0], ranges[:, 1])
peg_quat = np.array([1, 0, 0, 0])
peg_pose = np.concatenate([peg_position, peg_quat])
@@ -31,7 +35,7 @@ def sample_insertion_pose():
z_range = [0.05, 0.05]
ranges = np.vstack([x_range, y_range, z_range])
socket_position = np.random.uniform(ranges[:, 0], ranges[:, 1])
socket_position = rng.uniform(ranges[:, 0], ranges[:, 1])
socket_quat = np.array([1, 0, 0, 0])
socket_pose = np.concatenate([socket_position, socket_quat])

20
lerobot/common/envs/utils.py
View File

@@ -6,12 +6,20 @@ from lerobot.common.transforms import apply_inverse_transform
def preprocess_observation(observation, transform=None):
# map to expected inputs for the policy
obs = {
"observation.image": torch.from_numpy(observation["pixels"]).float(),
"observation.state": torch.from_numpy(observation["agent_pos"]).float(),
}
# convert to (b c h w) torch format
obs["observation.image"] = einops.rearrange(obs["observation.image"], "b h w c -> b c h w")
obs = {}
if isinstance(observation["pixels"], dict):
imgs = {f"observation.images.{key}": img for key, img in observation["pixels"].items()}
else:
imgs = {"observation.image": observation["pixels"]}
for imgkey, img in imgs.items():
img = torch.from_numpy(img).float()
# convert to (b c h w) torch format
img = einops.rearrange(img, "b h w c -> b c h w")
obs[imgkey] = img
obs["observation.state"] = torch.from_numpy(observation["agent_pos"]).float()
# apply same transforms as in training
if transform is not None: