-
Notifications
You must be signed in to change notification settings - Fork 641
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
WIP Aloha env tests pass Rendering works (fps look fast tho? TODO action bounding is too wide [-1,1]) Update README Copy past from act repo Remove download.py add a WIP for Simxarm Remove download.py add a WIP for Simxarm Add act yaml (TODO: try train.py) Training can runs (TODO: eval) Add tasks without end_effector that are compatible with dataset, Eval can run (TODO: training and pretrained model) Add AbstractEnv, Refactor AlohaEnv, Add rendering_hook in env, Minor modifications, (TODO: Refactor Pusht and Simxarm) poetry lock fix bug in compute_stats for action normalization fix more bugs in normalization fix training fix import PushtEnv inheriates AbstractEnv, Improve factory Normalization Add _make_env to EnvAbstract Add call_rendering_hooks to pusht env SimxarmEnv inherites from AbstractEnv (NOT TESTED) Add aloha tests artifacts + update pusht stats fix image normalization: before env was in [0,1] but dataset in [0,255], and now both in [0,255] Small fix on simxarm Add next to obs Add top camera to Aloha env (TODO: make it compatible with set of cameras) Add top camera to Aloha env (TODO: make it compatible with set of cameras)
- Loading branch information
Showing
116 changed files
with
3,658 additions
and
301 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,80 @@ | ||
import abc | ||
from collections import deque | ||
from typing import Optional | ||
|
||
from tensordict import TensorDict | ||
from torchrl.envs import EnvBase | ||
|
||
|
||
class AbstractEnv(EnvBase): | ||
def __init__( | ||
self, | ||
task, | ||
frame_skip: int = 1, | ||
from_pixels: bool = False, | ||
pixels_only: bool = False, | ||
image_size=None, | ||
seed=1337, | ||
device="cpu", | ||
num_prev_obs=1, | ||
num_prev_action=0, | ||
): | ||
super().__init__(device=device, batch_size=[]) | ||
self.task = task | ||
self.frame_skip = frame_skip | ||
self.from_pixels = from_pixels | ||
self.pixels_only = pixels_only | ||
self.image_size = image_size | ||
self.num_prev_obs = num_prev_obs | ||
self.num_prev_action = num_prev_action | ||
self._rendering_hooks = [] | ||
|
||
if pixels_only: | ||
assert from_pixels | ||
if from_pixels: | ||
assert image_size | ||
|
||
self._make_env() | ||
self._make_spec() | ||
self._current_seed = self.set_seed(seed) | ||
|
||
if self.num_prev_obs > 0: | ||
self._prev_obs_image_queue = deque(maxlen=self.num_prev_obs) | ||
self._prev_obs_state_queue = deque(maxlen=self.num_prev_obs) | ||
if self.num_prev_action > 0: | ||
raise NotImplementedError() | ||
# self._prev_action_queue = deque(maxlen=self.num_prev_action) | ||
|
||
def register_rendering_hook(self, func): | ||
self._rendering_hooks.append(func) | ||
|
||
def call_rendering_hooks(self): | ||
for func in self._rendering_hooks: | ||
func(self) | ||
|
||
def reset_rendering_hooks(self): | ||
self._rendering_hooks = [] | ||
|
||
@abc.abstractmethod | ||
def render(self, mode="rgb_array", width=640, height=480): | ||
raise NotImplementedError() | ||
|
||
@abc.abstractmethod | ||
def _reset(self, tensordict: Optional[TensorDict] = None): | ||
raise NotImplementedError() | ||
|
||
@abc.abstractmethod | ||
def _step(self, tensordict: TensorDict): | ||
raise NotImplementedError() | ||
|
||
@abc.abstractmethod | ||
def _make_env(self): | ||
raise NotImplementedError() | ||
|
||
@abc.abstractmethod | ||
def _make_spec(self): | ||
raise NotImplementedError() | ||
|
||
@abc.abstractmethod | ||
def _set_seed(self, seed: Optional[int]): | ||
raise NotImplementedError() |
59 changes: 59 additions & 0 deletions
59
lerobot/common/envs/aloha/assets/bimanual_viperx_end_effector_insertion.xml
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,59 @@ | ||
<mujoco> | ||
<include file="scene.xml"/> | ||
<include file="vx300s_dependencies.xml"/> | ||
|
||
<equality> | ||
<weld body1="mocap_left" body2="vx300s_left/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" /> | ||
<weld body1="mocap_right" body2="vx300s_right/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" /> | ||
</equality> | ||
|
||
|
||
<worldbody> | ||
<include file="vx300s_left.xml" /> | ||
<include file="vx300s_right.xml" /> | ||
|
||
<body mocap="true" name="mocap_left" pos="0.095 0.50 0.425"> | ||
<site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_left_site1" rgba="1 0 0 1"/> | ||
<site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_left_site2" rgba="1 0 0 1"/> | ||
<site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_left_site3" rgba="1 0 0 1"/> | ||
</body> | ||
<body mocap="true" name="mocap_right" pos="-0.095 0.50 0.425"> | ||
<site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_right_site1" rgba="1 0 0 1"/> | ||
<site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_right_site2" rgba="1 0 0 1"/> | ||
<site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_right_site3" rgba="1 0 0 1"/> | ||
</body> | ||
|
||
<body name="peg" pos="0.2 0.5 0.05"> | ||
<joint name="red_peg_joint" type="free" frictionloss="0.01" /> | ||
<inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" /> | ||
<geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg" rgba="1 0 0 1" /> | ||
</body> | ||
|
||
<body name="socket" pos="-0.2 0.5 0.05"> | ||
<joint name="blue_socket_joint" type="free" frictionloss="0.01" /> | ||
<inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" /> | ||
<!-- <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg_ref" rgba="1 0 0 1" />--> | ||
<geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 -0.02" size="0.06 0.018 0.002" type="box" name="socket-1" rgba="0 0 1 1" /> | ||
<geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 0.02" size="0.06 0.018 0.002" type="box" name="socket-2" rgba="0 0 1 1" /> | ||
<geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0.02 0" size="0.06 0.002 0.018" type="box" name="socket-3" rgba="0 0 1 1" /> | ||
<geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 -0.02 0" size="0.06 0.002 0.018" type="box" name="socket-4" rgba="0 0 1 1" /> | ||
<geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.04 0.01 0.01" type="box" name="pin" rgba="1 0 0 1" /> | ||
</body> | ||
|
||
</worldbody> | ||
|
||
<actuator> | ||
<position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_left/left_finger" kp="200" user="1"/> | ||
<position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_left/right_finger" kp="200" user="1"/> | ||
|
||
<position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_right/left_finger" kp="200" user="1"/> | ||
<position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_right/right_finger" kp="200" user="1"/> | ||
|
||
</actuator> | ||
|
||
<keyframe> | ||
<key qpos="0 -0.96 1.16 0 -0.3 0 0.024 -0.024 0 -0.96 1.16 0 -0.3 0 0.024 -0.024 0.2 0.5 0.05 1 0 0 0 -0.2 0.5 0.05 1 0 0 0"/> | ||
</keyframe> | ||
|
||
|
||
</mujoco> |
48 changes: 48 additions & 0 deletions
48
lerobot/common/envs/aloha/assets/bimanual_viperx_end_effector_transfer_cube.xml
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,48 @@ | ||
<mujoco> | ||
<include file="scene.xml"/> | ||
<include file="vx300s_dependencies.xml"/> | ||
|
||
<equality> | ||
<weld body1="mocap_left" body2="vx300s_left/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" /> | ||
<weld body1="mocap_right" body2="vx300s_right/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" /> | ||
</equality> | ||
|
||
|
||
<worldbody> | ||
<include file="vx300s_left.xml" /> | ||
<include file="vx300s_right.xml" /> | ||
|
||
<body mocap="true" name="mocap_left" pos="0.095 0.50 0.425"> | ||
<site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_left_site1" rgba="1 0 0 1"/> | ||
<site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_left_site2" rgba="1 0 0 1"/> | ||
<site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_left_site3" rgba="1 0 0 1"/> | ||
</body> | ||
<body mocap="true" name="mocap_right" pos="-0.095 0.50 0.425"> | ||
<site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_right_site1" rgba="1 0 0 1"/> | ||
<site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_right_site2" rgba="1 0 0 1"/> | ||
<site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_right_site3" rgba="1 0 0 1"/> | ||
</body> | ||
|
||
<body name="box" pos="0.2 0.5 0.05"> | ||
<joint name="red_box_joint" type="free" frictionloss="0.01" /> | ||
<inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" /> | ||
<geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.02 0.02 0.02" type="box" name="red_box" rgba="1 0 0 1" /> | ||
</body> | ||
|
||
</worldbody> | ||
|
||
<actuator> | ||
<position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_left/left_finger" kp="200" user="1"/> | ||
<position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_left/right_finger" kp="200" user="1"/> | ||
|
||
<position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_right/left_finger" kp="200" user="1"/> | ||
<position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_right/right_finger" kp="200" user="1"/> | ||
|
||
</actuator> | ||
|
||
<keyframe> | ||
<key qpos="0 -0.96 1.16 0 -0.3 0 0.024 -0.024 0 -0.96 1.16 0 -0.3 0 0.024 -0.024 0.2 0.5 0.05 1 0 0 0"/> | ||
</keyframe> | ||
|
||
|
||
</mujoco> |
Oops, something went wrong.