from rlbench import DomainRandomizationEnvironment
from rlbench import RandomizeEvery
from rlbench import VisualRandomizationConfig
from rlbench.action_modes import ActionMode, ArmActionMode
from rlbench.tasks import OpenDoor
import numpy as np
# We will borrow some from the tests dir
rand_config = VisualRandomizationConfig(
image_directory='../tests/unit/assets/textures')
action_mode = ActionMode(ArmActionMode.ABS_JOINT_VELOCITY)
env = DomainRandomizationEnvironment(action_mode,
randomize_every=RandomizeEvery.EPISODE,
frequency=1,
visual_randomization_config=None)
env.launch()
live_demos = True
task = env.get_task(OpenDoor)
demos = task.get_demos(2, live_demos=live_demos)
descriptions, obs = task.reset()
obs, reward, terminate = task.step(action)
print('Done')
env.shutdown()
pose[:3] += pos
pose[3:] = [
perturbed_quat_wxyz.x, perturbed_quat_wxyz.y,
perturbed_quat_wxyz.z, perturbed_quat_wxyz.w
]
obj_poses[name] = pose
return obj_poses
if __name__ == "__main__":
# action_mode = ActionMode(ArmActionMode.DELTA_EE_POSE) # See rlbench/action_modes.py for other action modes
# action_mode = ActionMode(ArmActionMode.DELTA_EE_POSE_PLAN)
action_mode = ActionMode(ArmActionMode.ABS_EE_POSE_PLAN)
env = Environment(action_mode, '', ObservationConfig(), False, frequency=5)
# task = env.get_task(StackBlocks) # available tasks: EmptyContainer, PlayJenga, PutGroceriesInCupboard, SetTheTable
task = env.get_task(
PlayJenga
) # available tasks: EmptyContainer, PlayJenga, PutGroceriesInCupboard, SetTheTable
obj_pose_sensor = NoisyObjectPoseSensor(env)
descriptions, obs = task.reset()
agent = Agent(obs, obj_pose_sensor.get_poses())
while True:
# Getting noisy object poses
obj_poses = obj_pose_sensor.get_poses()
# Getting various fields from obs
def __init__(self, task_class, observation_mode='state', randomization_mode="none",
rand_config=None, img_size=256, special_start=[], fixed_grip=-1, force_randomly_place=False, force_change_position=False, sparse=False, not_special_p = 0, ground_p = 0, special_is_grip=False, altview=False, procedural_ind=0, procedural_mode='same', procedural_set = [], is_mesh_obs=False, blank=False, COLOR_TUPLE=[1,0,0]):
# blank is 0s agent.
self.blank = blank
#altview is whether to have second camera angle or not. True/False, "both" to concatentae the observations.
self.altview=altview
self.img_size=img_size
self.sparse = sparse
self.task_class = task_class
self._observation_mode = observation_mode
self._randomization_mode = randomization_mode
#special start is a list of actions to take at the beginning.
self.special_start = special_start
#fixed_grip temp hack for keeping the gripper a certain way. Change later. 0 for fixed closed, 0.1 for fixed open, -1 for not fixed
self.fixed_grip = fixed_grip
#to force the task to be randomly placed
self.force_randomly_place = force_randomly_place
#force the task to change position in addition to rotation
self.force_change_position = force_change_position
obs_config = ObservationConfig()
if observation_mode == 'state':
obs_config.set_all_high_dim(False)
obs_config.set_all_low_dim(True)
elif observation_mode == 'vision' or observation_mode=="visiondepth" or observation_mode=="visiondepthmask":
# obs_config.set_all(True)
obs_config.set_all_high_dim(False)
obs_config.set_all_low_dim(True)
else:
raise ValueError(
'Unrecognised observation_mode: %s.' % observation_mode)
action_mode = ActionMode(ArmActionMode.DELTA_EE_POSE_PLAN_NOQ)
print("using delta pose pan")
if randomization_mode == "random":
objs = ['target', 'boundary', 'Floor', 'Roof', 'Wall1', 'Wall2', 'Wall3', 'Wall4', 'diningTable_visible']
if rand_config is None:
assert False
self.env = DomainRandomizationEnvironment(
action_mode, obs_config=obs_config, headless=True,
randomize_every=RandomizeEvery.EPISODE, frequency=1,
visual_randomization_config=rand_config
)
else:
self.env = Environment(
action_mode, obs_config=obs_config, headless=True
)
self.env.launch()
self.task = self.env.get_task(task_class)
# Probability. Currently used for probability that pick and lift task will start off gripper at a certain location (should probs be called non_special p)
self.task._task.not_special_p = not_special_p
# Probability that ground goal.
self.task._task.ground_p = ground_p
# For the "special" case, whether to grip the object or just hover above it.
self.task._task.special_is_grip = special_is_grip
# for procedural env
self.task._task.procedural_ind = procedural_ind
# procedural mode: same, increase, or random.
self.task._task.procedural_mode = procedural_mode
# ideally a list-like object, dictates the indices to sample from each episode.
self.task._task.procedural_set = procedural_set
# if state obs is mesh obs
self.task._task.is_mesh_obs = is_mesh_obs
self.task._task.sparse = sparse
self.task._task.COLOR_TUPLE = COLOR_TUPLE
_, obs = self.task.reset()
cam_placeholder = Dummy('cam_cinematic_placeholder')
cam_pose = cam_placeholder.get_pose().copy()
#custom pose
cam_pose = [ 1.59999931, 0. , 2.27999949 , 0.65328157, -0.65328145, -0.27059814, 0.27059814]
cam_pose[0] = 1
cam_pose[2] = 1.5
self.frontcam = VisionSensor.create([img_size, img_size])
self.frontcam.set_pose(cam_pose)
self.frontcam.set_render_mode(RenderMode.OPENGL)
self.frontcam.set_perspective_angle(60)
self.frontcam.set_explicit_handling(1)
self.frontcam_mask = VisionSensor.create([img_size, img_size])
self.frontcam_mask.set_pose(cam_pose)
self.frontcam_mask.set_render_mode(RenderMode.OPENGL_COLOR_CODED)
self.frontcam_mask.set_perspective_angle(60)
self.frontcam_mask.set_explicit_handling(1)
if altview:
alt_pose = [0.25 , -0.65 , 1.5, 0, 0.93879825 ,0.34169483 , 0]
self.altcam = VisionSensor.create([img_size, img_size])
self.altcam.set_pose(alt_pose)
self.altcam.set_render_mode(RenderMode.OPENGL)
self.altcam.set_perspective_angle(60)
self.altcam.set_explicit_handling(1)
self.altcam_mask = VisionSensor.create([img_size, img_size])
self.altcam_mask.set_pose(alt_pose)
self.altcam_mask.set_render_mode(RenderMode.OPENGL_COLOR_CODED)
self.altcam_mask.set_perspective_angle(60)
self.altcam_mask.set_explicit_handling(1)
self.action_space = spaces.Box(
low=-1.0, high=1.0, shape=(action_mode.action_size,))
if observation_mode == 'state':
self.observation_space = spaces.Dict({
"observation": spaces.Box(
low=-np.inf, high=np.inf, shape=self.task._task.get_state_obs().shape),
"achieved_goal": spaces.Box(
low=-np.inf, high=np.inf, shape=self.task._task.get_achieved_goal().shape),
'desired_goal': spaces.Box(
low=-np.inf, high=np.inf, shape=self.task._task.get_desired_goal().shape)
})
# Use the frontvision cam
elif observation_mode == 'vision':
self.frontcam.handle_explicitly()
self.observation_space = spaces.Dict({
"state": spaces.Box(
low=-np.inf, high=np.inf,
shape=obs.get_low_dim_data().shape),
"observation": spaces.Box(
low=0, high=1, shape=self.frontcam.capture_rgb().transpose(2,0,1).flatten().shape,
)
})
if altview == "both":
example = self.frontcam.capture_rgb().transpose(2,0,1).flatten()
self.observation_space = spaces.Dict({
"state": spaces.Box(
low=-np.inf, high=np.inf,
shape=obs.get_low_dim_data().shape),
"observation": spaces.Box(
low=0, high=1, shape=np.array([example,example]).shape,
)
})
elif observation_mode == 'visiondepth':
self.frontcam.handle_explicitly()
self.observation_space = spaces.Dict({
"state": spaces.Box(
low=-np.inf, high=np.inf,
shape=obs.get_low_dim_data().shape),
"observation": spaces.Box(
#thinking about not flattening the shape, because primarily used for dataset and not for training
low=0, high=1, shape=np.array([self.frontcam.capture_rgb().transpose(2,0,1), self.frontcam.capture_depth()[None,...]]).shape,
)
})
elif observation_mode == 'visiondepthmask':
self.frontcam.handle_explicitly()
self.frontcam_mask.handle_explicitly()
self.observation_space = spaces.Dict({
"state": spaces.Box(
low=-np.inf, high=np.inf,
shape=obs.get_low_dim_data().shape),
"observation": spaces.Box(
#thinking about not flattening the shape, because primarily used for dataset and not for training
low=0, high=1, shape=np.array([self.frontcam.capture_rgb().transpose(2,0,1), self.frontcam.capture_depth()[None,...], rgb_handles_to_mask(self.frontcam_mask.capture_rgb())]).shape,
)
})
if altview == "both":
self.observation_space = spaces.Dict({
"state": spaces.Box(
low=-np.inf, high=np.inf,
shape=obs.get_low_dim_data().shape),
"observation": spaces.Box(
#thinking about not flattening the shape, because primarily used for dataset and not for training
low=0, high=1, shape=np.array([self.frontcam.capture_rgb().transpose(2,0,1), self.frontcam.capture_rgb().transpose(2,0,1), self.frontcam.capture_depth()[None,...], rgb_handles_to_mask(self.frontcam_mask.capture_rgb())]).shape,
)
})
self._gym_cam = None
path_joints = np.hstack(
(path_points, np.zeros((path_points.shape[0], 1))))
i = 0
while not path._path_done and i < path_joints.shape[0]:
task.step(path_joints[i])
i += 1
if __name__ == "__main__":
# Initializes environment and task
mode = "abs_joint_pos" # ee_pose_plan
# mode = "ee_pose_plan"
if (mode == "ee_pose_plan"):
action_mode = ActionMode(
ArmActionMode.ABS_EE_POSE_PLAN
) # See rlbench/action_modes.py for other action modes
elif (mode == "abs_joint_pos"):
action_mode = ActionMode(ArmActionMode.ABS_JOINT_POSITION)
else:
print("Mode Not Found")
env = Environment(action_mode,
'',
ObservationConfig(),
False,
static_positions=False)
task = env.get_task(
PutGroceriesInCupboard
) # available tasks: EmptyContainer, PlayJenga, PutGroceriesInCupboard, SetTheTable
task.reset()
def run(i, lock, task_index, variation_count, results, file_lock, tasks):
"""Each thread will choose one task and variation, and then gather
all the episodes_per_task for that variation."""
# Initialise each thread with random seed
np.random.seed(None)
num_tasks = len(tasks)
img_size = list(map(int, FLAGS.image_size))
obs_config = ObservationConfig()
obs_config.set_all(True)
obs_config.right_shoulder_camera.image_size = img_size
obs_config.left_shoulder_camera.image_size = img_size
obs_config.wrist_camera.image_size = img_size
obs_config.front_camera.image_size = img_size
# We want to save the masks as rgb encodings.
obs_config.left_shoulder_camera.masks_as_one_channel = False
obs_config.right_shoulder_camera.masks_as_one_channel = False
obs_config.wrist_camera.masks_as_one_channel = False
obs_config.front_camera.masks_as_one_channel = False
if FLAGS.renderer == 'opengl':
obs_config.right_shoulder_camera.render_mode = RenderMode.OPENGL
obs_config.left_shoulder_camera.render_mode = RenderMode.OPENGL
obs_config.wrist_camera.render_mode = RenderMode.OPENGL
obs_config.front_camera.render_mode = RenderMode.OPENGL
rlbench_env = Environment(action_mode=ActionMode(),
obs_config=obs_config,
headless=True)
rlbench_env.launch()
task_env = None
tasks_with_problems = results[i] = ''
while True:
# Figure out what task/variation this thread is going to do
with lock:
if task_index.value >= num_tasks:
print('Process', i, 'finished')
break
my_variation_count = variation_count.value
t = tasks[task_index.value]
task_env = rlbench_env.get_task(t)
var_target = task_env.variation_count()
if FLAGS.variations >= 0:
var_target = np.minimum(FLAGS.variations, var_target)
if my_variation_count >= var_target:
# If we have reached the required number of variations for this
# task, then move on to the next task.
variation_count.value = my_variation_count = 0
task_index.value += 1
variation_count.value += 1
if task_index.value >= num_tasks:
print('Process', i, 'finished')
break
t = tasks[task_index.value]
print('Process', i, 'collecting task:', task_env.get_name(),
'// variation:', my_variation_count)
task_env = rlbench_env.get_task(t)
task_env.set_variation(my_variation_count)
obs, descriptions = task_env.reset()
variation_path = os.path.join(FLAGS.save_path, task_env.get_name(),
VARIATIONS_FOLDER % my_variation_count)
check_and_make(variation_path)
with open(os.path.join(variation_path, VARIATION_DESCRIPTIONS),
'wb') as f:
pickle.dump(descriptions, f)
episodes_path = os.path.join(variation_path, EPISODES_FOLDER)
check_and_make(episodes_path)
abort_variation = False
for ex_idx in range(FLAGS.episodes_per_task):
attempts = 10
while attempts > 0:
try:
# TODO: for now we do the explicit looping.
demo, = task_env.get_demos(amount=1, live_demos=True)
except Exception as e:
attempts -= 1
if attempts > 0:
continue
problem = (
'Process %d failed collecting task %s (variation: %d, '
'example: %d). Skipping this task/variation.\n%s\n' %
(i, task_env.get_name(), my_variation_count, ex_idx,
str(e)))
print(problem)
tasks_with_problems += problem
abort_variation = True
break
episode_path = os.path.join(episodes_path,
EPISODE_FOLDER % ex_idx)
with file_lock:
save_demo(demo, episode_path)
break
if abort_variation:
break
results[i] = tasks_with_problems
rlbench_env.shutdown()
# https://github.com/stepjam/RLBench/blob/master/rlbench/observation_config.py#L5
# The depth Values are also single channel. Where depth will be of dims (width,height)
# https://github.com/stepjam/PyRep/blob/4a61f6756c3827db66423409632358de312b97e4/pyrep/objects/vision_sensor.py#L128
image_types = [
'left_shoulder_rgb', # (width,height,channel)
'left_shoulder_depth', # Depth is in Black and White : (width,height)
'left_shoulder_mask', # Mask can be single channel
'right_shoulder_rgb',
'right_shoulder_depth', # Depth is in Black and White
'right_shoulder_mask', # Mask can be single channel
'wrist_rgb',
'wrist_depth', # Depth is in Black and White
'wrist_mask' # Mask can be single channel
]
DEFAULT_ACTION_MODE = ActionMode(ArmActionMode.ABS_JOINT_VELOCITY)
DEFAULT_TASK = ReachTarget
class SimulationEnvionment():
"""
This can be a parent class from which we can have multiple child classes that
can diversify for different tasks and deeper functions within the tasks.
"""
def __init__(self,\
action_mode=DEFAULT_ACTION_MODE,\
task=DEFAULT_TASK,\
dataset_root='',\
headless=True):
obs_config = ObservationConfig()
obs_config.set_all(True)
def __setup_action_mode(self):
arm_action_mode_type = self.cfg["Agent"]["ArmActionMode"]
arm_action_mode = None
gripper_action_mode_type = self.cfg["Agent"]["GripperActionMode"]
gripper_action_mode = None
if arm_action_mode_type == "ABS_JOINT_VELOCITY":
# Absolute arm joint velocities
arm_action_mode = ArmActionMode.ABS_JOINT_VELOCITY
elif arm_action_mode_type == "DELTA_JOINT_VELOCITY":
# Change in arm joint velocities
arm_action_mode = ArmActionMode.DELTA_JOINT_VELOCITY
elif arm_action_mode == "ABS_JOINT_POSITION":
# Absolute arm joint positions/angles (in radians)
arm_action_mode = ArmActionMode.ABS_JOINT_VELOCITY
elif arm_action_mode == "DELTA_JOINT_POSITION":
# Change in arm joint positions/angles (in radians)
arm_action_mode = ArmActionMode.DELTA_JOINT_POSITION
elif arm_action_mode == "ABS_JOINT_TORQUE":
# Absolute arm joint forces/torques
arm_action_mode = ArmActionMode.ABS_JOINT_TORQUE
elif arm_action_mode == "DELTA_JOINT_TORQUE":
# Change in arm joint forces/torques
arm_action_mode = ArmActionMode.DELTA_JOINT_TORQUE
elif arm_action_mode == "ABS_EE_VELOCITY":
# Absolute end-effector velocity (position (3) and quaternion (4))
arm_action_mode = ArmActionMode.ABS_EE_VELOCITY
elif arm_action_mode == "DELTA_EE_VELOCITY":
# Change in end-effector velocity (position (3) and quaternion (4))
arm_action_mode = ArmActionMode.DELTA_EE_VELOCITY
elif arm_action_mode == "ABS_EE_POSE":
# Absolute end-effector pose (position (3) and quaternion (4))
arm_action_mode = ArmActionMode.ABS_EE_POSE
elif arm_action_mode == "DELTA_EE_POSE":
# Change in end-effector pose (position (3) and quaternion (4))
arm_action_mode = ArmActionMode.DELTA_EE_POSE
elif arm_action_mode == "ABS_EE_POSE_PLAN":
# Absolute end-effector pose (position (3) and quaternion (4))
# But does path planning between these points
arm_action_mode = ArmActionMode.ABS_EE_POSE_PLAN
elif arm_action_mode == "DELTA_EE_POSE_PLAN":
# Change in end-effector pose (position (3) and quaternion (4))
# But does path planning between these points
arm_action_mode = ArmActionMode.DELTA_EE_POSE_PLAN
else:
raise ValueError(
"%s is not a supported arm action mode. Please check your config-file."
% arm_action_mode_type)
if gripper_action_mode_type == "OPEN_AMOUNT":
gripper_action_mode = GripperActionMode.OPEN_AMOUNT
else:
raise ValueError(
"%s is not a supported gripper action mode. Please check your config-file."
% gripper_action_mode_type)
return ActionMode(arm=arm_action_mode, gripper=gripper_action_mode)
perturbed_quat_wxyz = quat_wxyz * gt_quat_wxyz
pose[:3] += pos
pose[3:] = [
perturbed_quat_wxyz.x, perturbed_quat_wxyz.y,
perturbed_quat_wxyz.z, perturbed_quat_wxyz.w
]
obj_poses[name] = pose
return obj_poses
if __name__ == "__main__":
action_mode = ActionMode(
ArmActionMode.DELTA_EE_POSE
) # See rlbench/action_modes.py for other action modes
env = Environment(action_mode, '', ObservationConfig(), False)
task = env.get_task(
StackBlocks
) # available tasks: EmptyContainer, PlayJenga, PutGroceriesInCupboard, SetTheTable
agent = RandomAgent()
obj_pose_sensor = NoisyObjectPoseSensor(env)
descriptions, obs = task.reset()
print(descriptions)
while True:
# Getting noisy object poses
obj_poses = obj_pose_sensor.get_poses()
# Getting various fields from obs
def __init__(
self,
task_name: str,
state_type: list = 'state',
):
# render_mode=None):
"""
create RL Bench environment
:param task_name: task names can be found in rlbench.tasks
:param state_type: state or vision or a sub list of state_types list like ['left_shoulder_rgb']
"""
if state_type == 'state' or state_type == 'vision' or isinstance(
state_type, list):
self._state_type = state_type
else:
raise ValueError(
'State type value error, your value is {}'.format(state_type))
# self._render_mode = render_mode
self._render_mode = None
obs_config = ObservationConfig()
obs_config.set_all(True)
action_mode = ActionMode(ArmActionMode.ABS_JOINT_VELOCITY)
self.env = Environment(action_mode,
obs_config=obs_config,
headless=True)
self.env.launch()
try:
self.task = self.env.get_task(
getattr(sys.modules[__name__], task_name))
except:
raise NotImplementedError
_, obs = self.task.reset()
self.spec = Spec(task_name)
if self._state_type == 'state':
self.observation_space = spaces.Box(
low=-np.inf, high=np.inf, shape=obs.get_low_dim_data().shape)
elif self._state_type == 'vision':
space_dict = OrderedDict()
space_dict["state"] = spaces.Box(
low=-np.inf, high=np.inf, shape=obs.get_low_dim_data().shape)
for i in [
"left_shoulder_rgb", "right_shoulder_rgb", "wrist_rgb",
"front_rgb"
]:
space_dict[i] = spaces.Box(low=0,
high=1,
shape=getattr(obs, i).shape)
self.observation_space = spaces.Dict(space_dict)
else:
space_dict = OrderedDict()
for name in self._state_type:
if name.split('_')[-1] in ('rgb'):
space_dict[name] = spaces.Box(low=0,
high=1,
shape=getattr(obs,
name).shape)
elif name.split('_')[-1] in ('depth', 'mask'):
space_dict[name] = spaces.Box(low=0,
high=1,
shape=(128, 128, 3))
print(space_dict[name].shape)
else:
space_dict[name] = spaces.Box(low=-np.inf,
high=np.inf,
shape=getattr(obs,
name).shape)
self.observation_space = spaces.Dict(space_dict)
self.action_space = spaces.Box(low=-1.0,
high=1.0,
shape=(self.env.action_size, ),
dtype=np.float32)
# while not done:
# done = path.step()
# a = path.visualize()
# self.env._scene.step()
# return done
if __name__ == "__main__":
# Get grasp planner using GQCNN
grasp_planner = GraspPlanner(model="GQCNN-2.0")
# Get large container empty detector
large_container_detector = container_detector(model='large_container_detector_model.pth')
# Get small container empty detector
small_container_detector = container_detector(model='small_container_detector_model.pth')
# Set Action Mode, See rlbench/action_modes.py for other action modes
action_mode = ActionMode(ArmActionMode.ABS_JOINT_POSITION)
# Create grasp controller with initialized environment and task
grasp_controller = GraspController(action_mode, static_positions=True)
# Reset task
descriptions, obs = grasp_controller.reset()
# The camera intrinsic in RLBench
camera_intr = CameraIntrinsics(fx=893.738, fy=893.738, cx=516, cy=386, frame='world', height=772, width=1032)
# The translation between camera and gripper
# TODO: Change the whole logic into detecting the object using GQCNN
object_initial_poses = {}
while True:
camera_to_gripper_translation = [0.022, 0, 0.095]
while True:
objs = grasp_controller.get_objects(add_noise=True)
# Go to home position
if i_episode % 50 == 0:
agent.save()
print('Save Model\n\rEpisode {}\tAverage Score: {:.2f}'.format(
i_episode, np.mean(scores_deque)))
return scores, actor_losses, critic_losses
obs_config = ObservationConfig()
obs_config.set_all(False)
obs_config.joint_velocities = True
obs_config.joint_forces = True
obs_config.joint_positions = True
action_mode = ActionMode(ArmActionMode.ABS_JOINT_TORQUE)
env = RLBenchEnv(
"ReachTarget",
state_type_list=[
"joint_positions",
"joint_velocities",
# 'left_shoulder_rgb',
# 'right_shoulder_rgb',
'task_low_dim_state',
])
state = env.reset()
action_dim = env.action_space.shape[0]
state_space = env.observation_space
agent = Agent(state_space,
HIDDEN_SIZE,
import numpy as np
from natsort import natsorted
from rlbench.environment import Environment
from rlbench.observation_config import ObservationConfig
from rlbench.action_modes import ArmActionMode, ActionMode
from rlbench.tasks import ReachTarget
from data.rlbench_data_collection import collect_demos
obs_config = ObservationConfig()
obs_config.front_camera.rgb = True
obs_config.wrist_camera.rgb = True
obs_config.joint_forces = False
action_mode = ActionMode(ArmActionMode.ABS_EE_POSE_WORLD_FRAME)
# action_mode = ActionMode(ArmActionMode.ABS_JOINT_POSITION)
# action_mode = ActionMode(ArmActionMode.ABS_JOINT_VELOCITY)
env = Environment(action_mode,
obs_config=obs_config,
robot_configuration='ur3robotiq',
headless=False)
import time
time.sleep(3)
demos = collect_demos(env,
ReachTarget,
"hola",
n_iter_per_var=1,
n_demos_per_iter=3)
def __init__(self, task_class, n_features, load, n_movements):
"""
Learn the Movement from demonstration.
:param task_class: Task that we aim to learn
:param n_features: Number of RBF in n_features
:param load: Load from data
:param n_movements: how many movements do we want to learn
"""
frequency = 200
# To use 'saved' demos, set the path below, and set live_demos=False
live_demos = not load
DATASET = '' if live_demos else 'datasets'
obs_config = ObservationConfig()
obs_config.set_all_low_dim(True)
obs_config.set_all_high_dim(False)
self._task_class = task_class
action_mode = ActionMode(ArmActionMode.ABS_JOINT_POSITION)
group = Group("rlbench", ["d%d" % i for i in range(7)] + ["gripper"])
env = Environment(action_mode, DATASET, obs_config, headless=True)
env.launch()
task = env.get_task(task_class)
trajectories = []
states = []
lengths = []
print("Start Demo")
demos = task.get_demos(n_movements, live_demos=live_demos)
print("End Demo")
init = True
for demo in demos:
trajectory = NamedTrajectory(*group.refs)
t = 0
for ob in demo:
if t == 0:
if init:
print("State dim: %d" % ob.task_low_dim_state.shape[0])
init = False
states.append(ob.task_low_dim_state)
kwargs = {
"d%d" % i: ob.joint_positions[i]
for i in range(ob.joint_positions.shape[0])
}
kwargs["gripper"] = ob.gripper_open
trajectory.notify(duration=1 / frequency, **kwargs)
t += 1
lengths.append(t / 200.)
trajectories.append(trajectory)
space = ClassicSpace(group,
n_features=n_features,
regularization=1E-15)
z = np.linspace(-0.2, 1.2, 1000)
Phi = space.get_phi(z)
for i in range(n_features):
plt.plot(z, Phi[:, i])
plt.show()
parameters = np.array([
np.concatenate([
s,
np.array([l]),
LearnTrajectory(space, traj).get_block_params()
]) for s, l, traj in zip(states, lengths, trajectories)
])
np.save("parameters/%s_%d.npy" % (task.get_name(), n_features),
parameters)
env.shutdown()
