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

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])