# Copyright 2023, Junjia LIU, jjliu@mae.cuhk.edu.hk
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from isaacgym import gymtorch, gymapi
from isaacgym.torch_utils import *
from rofunc.learning.RofuncRL.tasks.isaacgymenv.base.vec_task import VecTask
from rofunc.utils.oslab.path import get_rofunc_path
[docs]class IKEABaseTask(VecTask):
def __init__(self, cfg, rl_device, sim_device, graphics_device_id, headless, virtual_screen_capture, force_render):
self.cfg = cfg
self.max_episode_length = self.cfg["env"]["episodeLength"]
self.action_scale = self.cfg["env"]["actionScale"]
self.start_position_noise = self.cfg["env"]["startPositionNoise"]
self.start_rotation_noise = self.cfg["env"]["startRotationNoise"]
self.num_props = self.cfg["env"]["numProps"]
self.aggregate_mode = self.cfg["env"]["aggregateMode"]
self.dof_vel_scale = self.cfg["env"]["dofVelocityScale"]
self.dist_reward_scale = self.cfg["env"]["distRewardScale"]
self.rot_reward_scale = self.cfg["env"]["rotRewardScale"]
self.around_handle_reward_scale = self.cfg["env"]["aroundHandleRewardScale"]
self.open_reward_scale = self.cfg["env"]["openRewardScale"]
self.finger_dist_reward_scale = self.cfg["env"]["fingerDistRewardScale"]
self.action_penalty_scale = self.cfg["env"]["actionPenaltyScale"]
self.debug_viz = self.cfg["env"]["enableDebugVis"]
self.up_axis = "z"
self.up_axis_idx = 2
self.distX_offset = 0.04
self.dt = 1 / 60.
# prop dimensions
self.prop_width = 0.08
self.prop_height = 0.08
self.prop_length = 0.08
self.prop_spacing = 0.09
# num_obs = 23
# num_acts = 9
#
# self.cfg["env"]["numObservations"] = 23
# self.cfg["env"]["numActions"] = 9
super().__init__(config=self.cfg, rl_device=rl_device, sim_device=sim_device,
graphics_device_id=graphics_device_id, headless=headless,
virtual_screen_capture=virtual_screen_capture, force_render=force_render)
# get gym GPU state tensors
actor_root_state_tensor = self.gym.acquire_actor_root_state_tensor(self.sim)
dof_state_tensor = self.gym.acquire_dof_state_tensor(self.sim)
rigid_body_tensor = self.gym.acquire_rigid_body_state_tensor(self.sim)
self.gym.refresh_actor_root_state_tensor(self.sim)
self.gym.refresh_dof_state_tensor(self.sim)
self.gym.refresh_rigid_body_state_tensor(self.sim)
# create some wrapper tensors for different slices
self.robot_default_dof_pos = to_torch([1.157, -1.066, -0.155, -2.239, -1.841, 1.003, 0.469, 0.035, 0.035],
device=self.device)
self.dof_state = gymtorch.wrap_tensor(dof_state_tensor)
self.robot_dof_state = self.dof_state.view(self.num_envs, -1, 2)[:, :self.num_robot_dofs]
self.robot_dof_state = self.robot_dof_state[..., 0]
self.robot_dof_state = self.robot_dof_state[..., 1]
self.furniture_dof_state = self.dof_state.view(self.num_envs, -1, 2)[:, self.num_robot_dofs:]
self.furniture_dof_pos = self.furniture_dof_state[..., 0]
self.furniture_dof_vel = self.furniture_dof_state[..., 1]
self.rigid_body_states = gymtorch.wrap_tensor(rigid_body_tensor).view(self.num_envs, -1, 13)
self.num_bodies = self.rigid_body_states.shape[1]
self.root_state_tensor = gymtorch.wrap_tensor(actor_root_state_tensor).view(self.num_envs, -1, 13)
if self.num_props > 0:
self.prop_states = self.root_state_tensor[:, 2:]
self.num_dofs = self.gym.get_sim_dof_count(self.sim) // self.num_envs
self.robot_dof_targets = torch.zeros((self.num_envs, self.num_dofs), dtype=torch.float, device=self.device)
self.global_indices = torch.arange(self.num_envs * (2 + self.num_props), dtype=torch.int32,
device=self.device).view(self.num_envs, -1)
self.reset_idx(torch.arange(self.num_envs, device=self.device))
[docs] def create_sim(self):
self.sim_params.up_axis = gymapi.UP_AXIS_Z
self.sim_params.gravity.x = 0
self.sim_params.gravity.y = 0
self.sim_params.gravity.z = -9.81
self.sim = super().create_sim(
self.device_id, self.graphics_device_id, self.physics_engine, self.sim_params)
self._create_ground_plane()
self._create_envs(self.num_envs, self.cfg["env"]['envSpacing'], int(np.sqrt(self.num_envs)))
def _create_ground_plane(self):
plane_params = gymapi.PlaneParams()
plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0)
self.gym.add_ground(self.sim, plane_params)
def _create_envs(self, num_envs, spacing, num_per_row):
lower = gymapi.Vec3(-spacing, -spacing, 0.0)
upper = gymapi.Vec3(spacing, spacing, spacing)
# get rofunc path from rofunc package metadata
rofunc_path = get_rofunc_path()
asset_root = os.path.join(rofunc_path, "simulator/assets")
robot_asset_file = "urdf/robot_description/robots/robot_panda.urdf"
furniture_asset_file = "urdf/objects/bed_dalselv_0270.xml"
# if "asset" in self.cfg["env"]:
# asset_root = os.path.join(os.path.dirname(os.path.abspath(__file__)), self.cfg["env"]["asset"].get("assetRoot", asset_root))
# robot_asset_file = self.cfg["env"]["asset"].get("assetFileNameFranka", robot_asset_file)
# furniture_asset_file = self.cfg["env"]["asset"].get("assetFileNameCabinet", furniture_asset_file)
# load franka asset
asset_options = gymapi.AssetOptions()
asset_options.flip_visual_attachments = True
asset_options.fix_base_link = True
asset_options.collapse_fixed_joints = True
asset_options.disable_gravity = True
asset_options.thickness = 0.001
asset_options.default_dof_drive_mode = gymapi.DOF_MODE_POS
asset_options.use_mesh_materials = True
robot_asset = self.gym.load_asset(self.sim, asset_root, robot_asset_file, asset_options)
robot_dof_stiffness = to_torch([400, 400, 400, 400, 400, 400, 400, 1.0e6, 1.0e6], dtype=torch.float,
device=self.device)
robot_dof_damping = to_torch([80, 80, 80, 80, 80, 80, 80, 1.0e2, 1.0e2], dtype=torch.float, device=self.device)
# load furniture asset
asset_options.flip_visual_attachments = False
asset_options.collapse_fixed_joints = True
asset_options.disable_gravity = False
asset_options.default_dof_drive_mode = gymapi.DOF_MODE_NONE
asset_options.armature = 0.005
furniture_asset = self.gym.load_asset(self.sim, asset_root, furniture_asset_file, asset_options)
self.num_robot_bodies = self.gym.get_asset_rigid_body_count(robot_asset)
self.num_robot_dofs = self.gym.get_asset_dof_count(robot_asset)
self.num_furniture_bodies = self.gym.get_asset_rigid_body_count(furniture_asset)
self.num_furniture_dofs = self.gym.get_asset_dof_count(furniture_asset)
print("num robot bodies: ", self.num_robot_bodies)
print("num robot dofs: ", self.num_robot_dofs)
print("num furniture bodies: ", self.num_furniture_bodies)
print("num furniture dofs: ", self.num_furniture_dofs)
# set franka dof properties
robot_dof_props = self.gym.get_asset_dof_properties(robot_asset)
self.robot_dof_lower_limits = []
self.robot_dof_upper_limits = []
for i in range(self.num_robot_dofs):
robot_dof_props['driveMode'][i] = gymapi.DOF_MODE_POS
if self.physics_engine == gymapi.SIM_PHYSX:
robot_dof_props['stiffness'][i] = robot_dof_stiffness[i]
robot_dof_props['damping'][i] = robot_dof_damping[i]
else:
robot_dof_props['stiffness'][i] = 7000.0
robot_dof_props['damping'][i] = 50.0
self.robot_dof_lower_limits.append(robot_dof_props['lower'][i])
self.robot_dof_upper_limits.append(robot_dof_props['upper'][i])
self.robot_dof_lower_limits = to_torch(self.robot_dof_lower_limits, device=self.device)
self.robot_dof_upper_limits = to_torch(self.robot_dof_upper_limits, device=self.device)
self.robot_dof_speed_scales = torch.ones_like(self.robot_dof_lower_limits)
self.robot_dof_speed_scales[[7, 8]] = 0.1
robot_dof_props['effort'][7] = 200
robot_dof_props['effort'][8] = 200
# set cabinet dof properties
furniture_dof_props = self.gym.get_asset_dof_properties(furniture_asset)
for i in range(self.num_furniture_dofs):
furniture_dof_props['damping'][i] = 10.0
# create prop assets
box_opts = gymapi.AssetOptions()
box_opts.density = 400
prop_asset = self.gym.create_box(self.sim, self.prop_width, self.prop_height, self.prop_width, box_opts)
robot_start_pose = gymapi.Transform()
robot_start_pose.p = gymapi.Vec3(1.0, 0.0, 0.0)
robot_start_pose.r = gymapi.Quat(0.0, 0.0, 1.0, 0.0)
furniture_start_pose = gymapi.Transform()
furniture_start_pose.p = gymapi.Vec3(*get_axis_params(0.4, self.up_axis_idx))
# compute aggregate size
num_robot_bodies = self.gym.get_asset_rigid_body_count(robot_asset)
num_robot_shapes = self.gym.get_asset_rigid_shape_count(robot_asset)
num_furniture_bodies = self.gym.get_asset_rigid_body_count(furniture_asset)
num_furniture_shapes = self.gym.get_asset_rigid_shape_count(furniture_asset)
num_prop_bodies = self.gym.get_asset_rigid_body_count(prop_asset)
num_prop_shapes = self.gym.get_asset_rigid_shape_count(prop_asset)
max_agg_bodies = num_robot_bodies + num_furniture_bodies + self.num_props * num_prop_bodies
max_agg_shapes = num_robot_shapes + num_furniture_shapes + self.num_props * num_prop_shapes
self.frankas = []
self.cabinets = []
self.default_prop_states = []
self.prop_start = []
self.envs = []
for i in range(self.num_envs):
# create env instance
env_ptr = self.gym.create_env(
self.sim, lower, upper, num_per_row
)
if self.aggregate_mode >= 3:
self.gym.begin_aggregate(env_ptr, max_agg_bodies, max_agg_shapes, True)
robot_actor = self.gym.create_actor(env_ptr, robot_asset, robot_start_pose, "franka", i, 1, 0)
self.gym.set_actor_dof_properties(env_ptr, robot_actor, robot_dof_props)
if self.aggregate_mode == 2:
self.gym.begin_aggregate(env_ptr, max_agg_bodies, max_agg_shapes, True)
furniture_pose = furniture_start_pose
furniture_pose.p.x += self.start_position_noise * (np.random.rand() - 0.5)
dz = 0.5 * np.random.rand()
dy = np.random.rand() - 0.5
furniture_pose.p.y += self.start_position_noise * dy
furniture_pose.p.z += self.start_position_noise * dz
furniture_actor = self.gym.create_actor(env_ptr, furniture_asset, furniture_pose, "cabinet", i, 2, 0)
self.gym.set_actor_dof_properties(env_ptr, furniture_actor, furniture_dof_props)
if self.aggregate_mode == 1:
self.gym.begin_aggregate(env_ptr, max_agg_bodies, max_agg_shapes, True)
if self.num_props > 0:
self.prop_start.append(self.gym.get_sim_actor_count(self.sim))
drawer_handle = self.gym.find_actor_rigid_body_handle(env_ptr, furniture_actor, "drawer_top")
drawer_pose = self.gym.get_rigid_transform(env_ptr, drawer_handle)
props_per_row = int(np.ceil(np.sqrt(self.num_props)))
xmin = -0.5 * self.prop_spacing * (props_per_row - 1)
yzmin = -0.5 * self.prop_spacing * (props_per_row - 1)
prop_count = 0
for j in range(props_per_row):
prop_up = yzmin + j * self.prop_spacing
for k in range(props_per_row):
if prop_count >= self.num_props:
break
propx = xmin + k * self.prop_spacing
prop_state_pose = gymapi.Transform()
prop_state_pose.p.x = drawer_pose.p.x + propx
propz, propy = 0, prop_up
prop_state_pose.p.y = drawer_pose.p.y + propy
prop_state_pose.p.z = drawer_pose.p.z + propz
prop_state_pose.r = gymapi.Quat(0, 0, 0, 1)
prop_handle = self.gym.create_actor(env_ptr, prop_asset, prop_state_pose,
"prop{}".format(prop_count), i, 0, 0)
prop_count += 1
prop_idx = j * props_per_row + k
self.default_prop_states.append([prop_state_pose.p.x, prop_state_pose.p.y, prop_state_pose.p.z,
prop_state_pose.r.x, prop_state_pose.r.y, prop_state_pose.r.z,
prop_state_pose.r.w,
0, 0, 0, 0, 0, 0])
if self.aggregate_mode > 0:
self.gym.end_aggregate(env_ptr)
self.envs.append(env_ptr)
self.frankas.append(robot_actor)
self.cabinets.append(furniture_actor)
self.hand_handle = self.gym.find_actor_rigid_body_handle(env_ptr, robot_actor, "panda_link7")
self.drawer_handle = self.gym.find_actor_rigid_body_handle(env_ptr, furniture_actor, "drawer_top")
self.lfinger_handle = self.gym.find_actor_rigid_body_handle(env_ptr, robot_actor, "panda_leftfinger")
self.rfinger_handle = self.gym.find_actor_rigid_body_handle(env_ptr, robot_actor, "panda_rightfinger")
self.default_prop_states = to_torch(self.default_prop_states, device=self.device, dtype=torch.float).view(
self.num_envs, self.num_props, 13)
self.init_data()
