Skip to content

Commit 58481b2

Browse files
KzZhengKzZheng
committed
fix bugs
1 parent da1a4a6 commit 58481b2

17 files changed

+298
-14
lines changed

‎README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ Users can use AMSolver to run the current tasks in the VLMbench or build new tas
1313
pip install -r requirements.txt
1414
pip install .
1515
```
16+
(meshlab for import new model; copy new simAddOnScript_PyRep.lua to Coppeliasim)
1617

1718
In the vlm folder, we have predefined some task categories and instance tasks for VLMbench. If you want to customize your own task, the scripts in the tools folder can be helpful.
1819

‎amsolver/gym/__init__.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
from gym.envs.registration import register
2+
import rlbench.backend.task as task
3+
import os
4+
from rlbench.utils import name_to_task_class
5+
from rlbench.gym.rlbench_env import RLBenchEnv
6+
7+
TASKS = [t for t in os.listdir(task.TASKS_PATH)
8+
if t != '__init__.py' and t.endswith('.py')]
9+
10+
for task_file in TASKS:
11+
task_name = task_file.split('.py')[0]
12+
task_class = name_to_task_class(task_name)
13+
register(
14+
id='%s-state-v0' % task_name,
15+
entry_point='rlbench.gym:RLBenchEnv',
16+
kwargs={
17+
'task_class': task_class,
18+
'observation_mode': 'state'
19+
}
20+
)
21+
register(
22+
id='%s-vision-v0' % task_name,
23+
entry_point='rlbench.gym:RLBenchEnv',
24+
kwargs={
25+
'task_class': task_class,
26+
'observation_mode': 'vision'
27+
}
28+
)

‎amsolver/gym/rlbench_env.py

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
from typing import Union, Dict, Tuple
2+
3+
import gym
4+
from gym import spaces
5+
from pyrep.const import RenderMode
6+
from pyrep.objects.dummy import Dummy
7+
from pyrep.objects.vision_sensor import VisionSensor
8+
from rlbench.environment import Environment
9+
from rlbench.action_modes import ArmActionMode, ActionMode
10+
from rlbench.observation_config import ObservationConfig
11+
import numpy as np
12+
13+
14+
class RLBenchEnv(gym.Env):
15+
"""An gym wrapper for RLBench."""
16+
17+
metadata = {'render.modes': ['human', 'rgb_array']}
18+
19+
def __init__(self, task_class, observation_mode='state',
20+
render_mode: Union[None, str] = None):
21+
self._observation_mode = observation_mode
22+
self._render_mode = render_mode
23+
obs_config = ObservationConfig()
24+
if observation_mode == 'state':
25+
obs_config.set_all_high_dim(False)
26+
obs_config.set_all_low_dim(True)
27+
elif observation_mode == 'vision':
28+
obs_config.set_all(True)
29+
else:
30+
raise ValueError(
31+
'Unrecognised observation_mode: %s.' % observation_mode)
32+
action_mode = ActionMode(ArmActionMode.ABS_JOINT_VELOCITY)
33+
self.env = Environment(
34+
action_mode, obs_config=obs_config, headless=True)
35+
self.env.launch()
36+
self.task = self.env.get_task(task_class)
37+
38+
_, obs = self.task.reset()
39+
40+
self.action_space = spaces.Box(
41+
low=-1.0, high=1.0, shape=(self.env.action_size,))
42+
43+
if observation_mode == 'state':
44+
self.observation_space = spaces.Box(
45+
low=-np.inf, high=np.inf, shape=obs.get_low_dim_data().shape)
46+
elif observation_mode == 'vision':
47+
self.observation_space = spaces.Dict({
48+
"state": spaces.Box(
49+
low=-np.inf, high=np.inf,
50+
shape=obs.get_low_dim_data().shape),
51+
"left_shoulder_rgb": spaces.Box(
52+
low=0, high=1, shape=obs.left_shoulder_rgb.shape),
53+
"right_shoulder_rgb": spaces.Box(
54+
low=0, high=1, shape=obs.right_shoulder_rgb.shape),
55+
"wrist_rgb": spaces.Box(
56+
low=0, high=1, shape=obs.wrist_rgb.shape),
57+
"front_rgb": spaces.Box(
58+
low=0, high=1, shape=obs.front_rgb.shape),
59+
})
60+
61+
if render_mode is not None:
62+
# Add the camera to the scene
63+
cam_placeholder = Dummy('cam_cinematic_placeholder')
64+
self._gym_cam = VisionSensor.create([640, 360])
65+
self._gym_cam.set_pose(cam_placeholder.get_pose())
66+
if render_mode == 'human':
67+
self._gym_cam.set_render_mode(RenderMode.OPENGL3_WINDOWED)
68+
else:
69+
self._gym_cam.set_render_mode(RenderMode.OPENGL3)
70+
71+
def _extract_obs(self, obs) -> Dict[str, np.ndarray]:
72+
if self._observation_mode == 'state':
73+
return obs.get_low_dim_data()
74+
elif self._observation_mode == 'vision':
75+
return {
76+
"state": obs.get_low_dim_data(),
77+
"left_shoulder_rgb": obs.left_shoulder_rgb,
78+
"right_shoulder_rgb": obs.right_shoulder_rgb,
79+
"wrist_rgb": obs.wrist_rgb,
80+
"front_rgb": obs.front_rgb,
81+
}
82+
83+
def render(self, mode='human') -> Union[None, np.ndarray]:
84+
if mode != self._render_mode:
85+
raise ValueError(
86+
'The render mode must match the render mode selected in the '
87+
'constructor. \nI.e. if you want "human" render mode, then '
88+
'create the env by calling: '
89+
'gym.make("reach_target-state-v0", render_mode="human").\n'
90+
'You passed in mode %s, but expected %s.' % (
91+
mode, self._render_mode))
92+
if mode == 'rgb_array':
93+
return self._gym_cam.capture_rgb()
94+
95+
def reset(self) -> Dict[str, np.ndarray]:
96+
descriptions, obs = self.task.reset()
97+
del descriptions # Not used.
98+
return self._extract_obs(obs)
99+
100+
def step(self, action) -> Tuple[Dict[str, np.ndarray], float, bool, dict]:
101+
obs, reward, terminate = self.task.step(action)
102+
return self._extract_obs(obs), reward, terminate, {}
103+
104+
def close(self) -> None:
105+
self.env.shutdown()

‎amsolver/robot_ttms/jaco.ttm

575 KB
Binary file not shown.

‎amsolver/robot_ttms/mico.ttm

576 KB
Binary file not shown.

‎amsolver/robot_ttms/panda.ttm

463 KB
Binary file not shown.

‎amsolver/robot_ttms/sawyer.ttm

532 KB
Binary file not shown.

‎amsolver/robot_ttms/ur5.ttm

345 KB
Binary file not shown.

‎requirements.txt

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,4 +5,6 @@ html-testRunner
55
setuptools
66
natsort
77
torch
8-
opencv-python
8+
opencv-python
9+
scipy
10+
open3d

‎setup.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,11 +31,11 @@ def get_version(rel_path):
3131
packages=[
3232
'amsolver',
3333
'amsolver.backend',
34-
'amsolver.tasks',
35-
'amsolver.task_ttms',
34+
# 'amsolver.tasks',
35+
# 'amsolver.task_ttms',
3636
'amsolver.robot_ttms',
3737
'amsolver.sim2real',
38-
'amsolver.assets',
38+
# 'amsolver.assets',
3939
'amsolver.gym'
4040
],
4141
package_data={'': ['*.ttm', '*.obj', '**/**/*.ttm', '**/**/*.obj'],

‎tools/dataset_generator_NLP.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,20 +25,20 @@
2525
FLAGS = flags.FLAGS
2626

2727
flags.DEFINE_string('save_path',
28-
'../vlmbench/train',
28+
'/home/kz-lab/Documents/research/dataset/VLMbench/val/unseen',
2929
'Where to save the demos.')
30-
flags.DEFINE_list('tasks', ['pour_demo_color'],
30+
flags.DEFINE_list('tasks', ['pick_cube_color'],
3131
'The tasks to collect. If empty, all tasks are collected.')
3232
flags.DEFINE_list('image_size', [224, 224],
3333
'The size of the images tp save.')
3434
flags.DEFINE_enum('renderer', 'opengl', ['opengl', 'opengl3'],
3535
'The renderer to use. opengl does not include shadows, '
3636
'but is faster.')
37-
flags.DEFINE_integer('processes', 64,
37+
flags.DEFINE_integer('processes', 4,
3838
'The number of parallel processes during collection.')
39-
flags.DEFINE_integer('episodes_per_task', 20,
39+
flags.DEFINE_integer('episodes_per_task', 5,
4040
'The number of episodes to collect per task.')
41-
flags.DEFINE_integer('variations', -1,
41+
flags.DEFINE_integer('variations', 1,
4242
'Number of variations to collect per task. -1 for all.')
4343

4444

‎tools/task_validator.py

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
from amsolver.backend.task import Task
2+
from amsolver.backend.scene import DemoError
3+
from amsolver.observation_config import ObservationConfig
4+
from pyrep import PyRep
5+
from pyrep.robots.arms.panda import Panda
6+
from pyrep.robots.end_effectors.panda_gripper import PandaGripper
7+
from amsolver.backend.const import TTT_FILE
8+
from amsolver.backend.scene import Scene
9+
from amsolver.backend.utils import task_file_to_task_class
10+
from amsolver.backend.task import TASKS_PATH
11+
from amsolver.backend.robot import Robot
12+
import numpy as np
13+
import os
14+
import argparse
15+
16+
DEMO_ATTEMPTS = 5
17+
MAX_VARIATIONS = 100
18+
19+
20+
class TaskValidationError(Exception):
21+
pass
22+
23+
24+
def task_smoke(task: Task, scene: Scene, variation=-1, demos=4, success=0.50,
25+
max_variations=3, test_demos=True):
26+
# -1 variations for all.
27+
28+
print('Running task validator on task: %s' % task.get_name())
29+
30+
# Loading
31+
scene.load(task)
32+
33+
# Number of variations
34+
variation_count = task.variation_count()
35+
if variation_count < 0:
36+
raise TaskValidationError(
37+
"The method 'variation_count' should return a number > 0.")
38+
39+
if variation_count > MAX_VARIATIONS:
40+
raise TaskValidationError(
41+
"This task had %d variations. Currently the limit is set to %d" %
42+
(variation_count, MAX_VARIATIONS))
43+
44+
# Base rotation bounds
45+
base_pos, base_ori = task.base_rotation_bounds()
46+
if len(base_pos) != 3 or len(base_ori) != 3:
47+
raise TaskValidationError(
48+
"The method 'base_rotation_bounds' should return a tuple "
49+
"containing a list of floats.")
50+
51+
# Boundary root
52+
root = task.boundary_root()
53+
if not root.still_exists():
54+
raise TaskValidationError(
55+
"The method 'boundary_root' should return a Dummy that is the root "
56+
"of the task.")
57+
58+
def variation_smoke(i):
59+
60+
print('Running task validator on variation: %d' % i)
61+
62+
attempt_result = False
63+
failed_demos = 0
64+
for j in range(DEMO_ATTEMPTS):
65+
failed_demos = run_demos(i)
66+
attempt_result = (failed_demos / float(demos) <= 1. - success)
67+
if attempt_result:
68+
break
69+
else:
70+
print('Failed on attempt %d. Trying again...' % j)
71+
72+
# Make sure we don't fail too often
73+
if not attempt_result:
74+
raise TaskValidationError(
75+
"Too many failed demo runs. %d of %d demos failed." % (
76+
failed_demos, demos))
77+
else:
78+
print('Variation %d of task %s is good!' % (i, task.get_name()))
79+
if test_demos:
80+
print('%d of %d demos were successful.' % (
81+
demos - failed_demos, demos))
82+
83+
def run_demos(variation_num):
84+
fails = 0
85+
for dr in range(demos):
86+
try:
87+
scene.reset()
88+
desc = scene.init_episode(variation_num, max_attempts=10)
89+
if not isinstance(desc, list) or len(desc) <= 0:
90+
raise TaskValidationError(
91+
"The method 'init_variation' should return a list of "
92+
"string descriptions.")
93+
if test_demos:
94+
demo = scene.get_demo(record=True)
95+
assert len(demo) > 0
96+
except DemoError as e:
97+
fails += 1
98+
print(e)
99+
except Exception as e:
100+
# TODO: check that we don't fall through all of these cases
101+
fails += 1
102+
print(e)
103+
104+
return fails
105+
106+
variations_to_test = [variation]
107+
if variation < 0:
108+
variations_to_test = list(range(
109+
np.minimum(variation_count, max_variations)))
110+
111+
# Task set-up
112+
scene.init_task()
113+
[variation_smoke(i) for i in variations_to_test]
114+
115+
116+
if __name__ == '__main__':
117+
parser = argparse.ArgumentParser()
118+
parser.add_argument("task", help="The task file to test.")
119+
args = parser.parse_args()
120+
121+
python_file = os.path.join(TASKS_PATH, args.task)
122+
if not os.path.isfile(python_file):
123+
raise RuntimeError('Could not find the task file: %s' % python_file)
124+
125+
task_class = task_file_to_task_class(args.task)
126+
127+
DIR_PATH = os.path.dirname(os.path.abspath(__file__))
128+
sim = PyRep()
129+
ttt_file = os.path.join(
130+
DIR_PATH, '..', 'rlbench', TTT_FILE)
131+
sim.launch(ttt_file, headless=True)
132+
sim.step_ui()
133+
sim.set_simulation_timestep(0.005)
134+
sim.step_ui()
135+
sim.start()
136+
137+
robot = Robot(Panda(), PandaGripper())
138+
139+
active_task = task_class(sim, robot)
140+
obs = ObservationConfig()
141+
obs.set_all(False)
142+
scene = Scene(sim, robot, obs)
143+
try:
144+
task_smoke(active_task, scene, variation=2)
145+
except TaskValidationError as e:
146+
sim.shutdown()
147+
raise e
148+
sim.shutdown()
149+
print('Validation successful!')

‎vlm/tasks/open_door.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
from amsolver.backend.conditions import DetectedCondition
1212
from amsolver.backend.spawn_boundary import SpawnBoundary
1313
from pyrep.objects.joint import Joint
14-
from amsolver.backend.utils import GenerateGraspPoints, Task_step, find_trajectory_from_success_grasps, get_local_grasp_pose, get_sorted_grasp_pose
1514

1615
door_states = {
1716
"open":["Fully close", "Slightly close"],

‎vlm/tasks/stack_cubes_color.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from amsolver.backend.conditions import DetectedCondition
1010
from amsolver.backend.spawn_boundary import SpawnBoundary
1111
from amsolver.backend.task import Task
12-
from amsolver.backend.utils import GenerateGraspPoints, Task_step, find_trajectory_from_success_grasps, get_local_grasp_pose, get_sorted_grasp_pose, scale_object, select_color
12+
from amsolver.backend.utils import scale_object, select_color
1313
from vlm.tasks.stack_cubes import StackCubes
1414
from pyrep.const import ObjectType, PrimitiveShape
1515

‎vlm/tasks/stack_cubes_relative.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from amsolver.backend.conditions import DetectedCondition
1010
from amsolver.backend.spawn_boundary import SpawnBoundary
1111
from amsolver.backend.task import Task
12-
from amsolver.backend.utils import GenerateGraspPoints, Task_step, find_trajectory_from_success_grasps, get_local_grasp_pose, get_relative_position_xy, get_sorted_grasp_pose, scale_object, select_color
12+
from amsolver.backend.utils import get_relative_position_xy, get_sorted_grasp_pose, scale_object, select_color
1313
from vlm.tasks.stack_cubes import StackCubes
1414
from pyrep.const import ObjectType, PrimitiveShape
1515

‎vlm/tasks/stack_cubes_shape.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from amsolver.backend.conditions import DetectedCondition
1010
from amsolver.backend.spawn_boundary import SpawnBoundary
1111
from amsolver.backend.task import Task
12-
from amsolver.backend.utils import GenerateGraspPoints, Task_step, find_trajectory_from_success_grasps, get_local_grasp_pose, get_sorted_grasp_pose, scale_object, select_color
12+
from amsolver.backend.utils import scale_object, select_color
1313
from vlm.tasks.stack_cubes import StackCubes
1414
from pyrep.const import ObjectType, PrimitiveShape
1515

‎vlm/tasks/stack_cubes_size.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from amsolver.backend.conditions import DetectedCondition
1010
from amsolver.backend.spawn_boundary import SpawnBoundary
1111
from amsolver.backend.task import Task
12-
from amsolver.backend.utils import GenerateGraspPoints, Task_step, find_trajectory_from_success_grasps, get_local_grasp_pose, get_sorted_grasp_pose, scale_object, select_color
12+
from amsolver.backend.utils import scale_object, select_color
1313
from vlm.tasks.stack_cubes import StackCubes
1414
from pyrep.const import ObjectType, PrimitiveShape
1515

0 commit comments

Comments
 (0)