def test(): from mujoco_py import load_model_from_xml, MjSim, MjViewer from mujoco_py.modder import TextureModder l, w, h, t, left, m = sample_cabinet() cab = build_cabinet(l, w, h, t, left) filename = 'test.urdf' with open(filename, "w") as text_file: text_file.write(cab.xml) # print(cab.xml) model = load_model_from_xml(cab.xml) sim = MjSim(model) viewer = MjViewer(sim) modder = TextureModder(sim) for name in sim.model.geom_names: modder.rand_all(name) t = 0 if cab.geom[3] == 1: sim.data.ctrl[0] = -0.2 else: sim.data.ctrl[0] = 0.2 while t < 1000: # for name in sim.model.geom_names: # modder.rand_all(name) sim.step() viewer.render() t += 1
def test(): import cv2 from mujoco_py import load_model_from_xml, MjSim, MjViewer from mujoco_py.modder import TextureModder for i in range(100): l, w, h, t, left, m = sample_refrigerator(False) fridge = build_refrigerator(l, w, h, t, left, set_pose=(3.0, 0.0, -1.0), set_rot=(0, 0, 0, 1)) model = load_model_from_xml(fridge.xml) sim = MjSim(model) viewer = MjViewer(sim) modder = TextureModder(sim) set_two_door_control(sim, 'refrigerator') t = 0 while t < 2000: sim.step() viewer.render() t += 1
def drop_cube_on_body_demo(): world = MujocoWorldBase() arena = EmptyArena() arena.set_origin([0, 0, 0]) world.merge(arena) soft_torso = SoftTorsoObject() obj = soft_torso.get_collision() box = BoxObject() box_obj = box.get_collision() obj.append(new_joint(name='soft_torso_free_joint', type='free')) box_obj.append(new_joint(name='box_free_joint', type='free')) world.merge_asset(soft_torso) world.worldbody.append(obj) world.worldbody.append(box_obj) # Place torso on ground collision_soft_torso = world.worldbody.find("./body") collision_soft_torso.set("pos", array_to_string(np.array([-0.1, 0, 0.1]))) model = world.get_model(mode="mujoco_py") sim = MjSim(model) viewer = MjViewer(sim) for _ in range(10000): sim.step() viewer.render()
def main(): parent_dir_path = str(pathlib.Path(__file__).parent.absolute()) fname = parent_dir_path + '/kinova_j2s6s300/j2s6s300.xml' model = load_model_from_path(fname) sim = MjSim(model) viewer = MjViewer(sim) t = 0 while True: if t == 1000: ndof = len(sim.data.qpos) captured_state = copy.deepcopy(sim.data.qpos) desired_vel = [0] * len(captured_state) kv = np.eye(ndof) * 10 if t < 1000: sim.data.ctrl[:] = sim.data.qfrc_bias[:] else: sim.data.ctrl[:] = mjc.pd(None, desired_vel, captured_state, sim, ndof=len(captured_state), kv=kv) sim.forward() sim.step() viewer.render() t += 1
class Environment: def __init__(self, sim): self.sim = sim self.viewer = None self.initial_state = self.sim.get_state() def is_done(self): x_pos = self.sim.get_state().qpos[0] angle = degrees(self.sim.get_state().qpos[1]) if (12.5 > angle > -12.5) and (0.99 > x_pos > -0.99): return False else: return True def get_reward(self): if not self.is_done(): return 1 else: return -1 def render(self): if not self.viewer: self.viewer = MjViewer(self.sim) self.viewer.render() def reset(self): random_pos = np.random.uniform(0., 0.05, size=(2)) random_vel = np.random.uniform(0., 0.05, size=(2)) self.initial_state = MjSimState(time=0.0, qpos=random_pos, qvel=random_vel, act=None, udd_state={}) self.sim.set_state(self.initial_state) return np.array([self.sim.get_state().qpos.tolist() + self.sim.get_state().qvel.tolist()])
def test(k=0): from mujoco_py import load_model_from_xml, MjSim, MjViewer from mujoco_py.modder import TextureModder l, w, h, t, left, m = sample_cabinet2() cab = build_cabinet2(l, w, h, t, left) # print(cab.xml) model = load_model_from_xml(cab.xml) sim = MjSim(model) viewer = MjViewer(sim) modder = TextureModder(sim) for name in sim.model.geom_names: modder.rand_all(name) set_two_door_control(sim) q1s = [] q2s = [] t = 0 # mode 0 indicates starting lc while t < 4000: # for name in sim.model.geom_names: # modder.rand_all(name) viewer.render() if t % 250 == 0: q1 = sim.data.qpos[0] q2 = sim.data.qpos[1] print(sim.data.qpos) q1s.append(q1) q2s.append(q2) sim.step() t += 1 # print(q1s) np.save('devdata/q1_' + str(k).zfill(3), q1s) np.save('devdata/q2_' + str(k).zfill(3), q2s)
def play(self, mocap_filepath): from mujoco_py import load_model_from_xml, MjSim, MjViewer curr_path = getcwd() xmlpath = '/mujoco/humanoid_deepmimic/envs/asset/dp_env_v2.xml' with open(curr_path + xmlpath) as fin: MODEL_XML = fin.read() model = load_model_from_xml(MODEL_XML) sim = MjSim(model) viewer = MjViewer(sim) self.read_raw_data(mocap_filepath) self.convert_raw_data() from time import sleep phase_offset = np.array([0.0, 0.0, 0.0]) while True: for k in range(len(self.data)): tmp_val = self.data_config[k] sim_state = sim.get_state() sim_state.qpos[:] = tmp_val[:] sim_state.qpos[:3] += phase_offset[:] sim.set_state(sim_state) sim.forward() viewer.render() sim_state = sim.get_state() phase_offset = sim_state.qpos[:3] phase_offset[2] = 0
def run_activations_fcn(est_activations, model_ver=0, timestep=0.005, Mj_render=False): #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! the q0 is now the chasis pos. needs to be fixed """ this function runs the predicted activations generatred from running the inverse map on the desired task kinematics """ #loading data #print("loading data") #task_kinematics=np.load("task_kinematics.npy") #est_task_activations=np.load("est_task_activations.npy") model = load_model_from_path( "./models/nmi_leg_w_chassis_v{}.xml".format(model_ver)) sim = MjSim(model) if Mj_render: viewer = MjViewer(sim) # viewer.cam.fixedcamid += 1 # viewer.cam.type = const.CAMERA_FIXED sim_state = sim.get_state() control_vector_length = sim.data.ctrl.__len__() print("control_vector_length: " + str(control_vector_length)) number_of_task_samples = est_activations.shape[0] real_attempt_positions = np.zeros((number_of_task_samples, 2)) real_attempt_activations = np.zeros((number_of_task_samples, 3)) chassis_pos = np.zeros(number_of_task_samples, ) sim.set_state(sim_state) for ii in range(number_of_task_samples): sim.data.ctrl[:] = est_activations[ii, :] sim.step() current_positions_array = sim.data.qpos[-2:] # current_kinematics_array=np.array( # [sim.data.qpos[0], # sim.data.qvel[0], # sim.data.qacc[0], # sim.data.qpos[1], # sim.data.qvel[1], # sim.data.qacc[1]] # ) chassis_pos[ii] = sim.data.get_geom_xpos("Chassis_frame")[0] real_attempt_positions[ii, :] = current_positions_array real_attempt_activations[ii, :] = sim.data.ctrl if Mj_render: viewer.render() real_attempt_kinematics = positions_to_kinematics_fcn( real_attempt_positions[:, 0], real_attempt_positions[:, 1], timestep=0.005) return real_attempt_kinematics, real_attempt_activations, chassis_pos
def run_activations_fcn(MuJoCo_model_name, est_activations, timestep=0.01, Mj_render=False): #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! the q0 is now the chasis pos. needs to be fixed """ this function runs the predicted activations generatred from running the inverse map on the desired task kinematics """ #loading data #print("loading data") #task_kinematics=np.load("task_kinematics.npy") #est_task_activations=np.load("est_task_activations.npy") MuJoCo_model = load_model_from_path("./models/"+MuJoCo_model_name) sim = MjSim(MuJoCo_model) if Mj_render: viewer = MjViewer(sim) # to move it to the mounted camera viewer.cam.fixedcamid += 1 viewer.cam.type = const.CAMERA_FIXED # # to record the video # viewer._record_video = True # # recording path # viewer._video_path = "~/Documents/"+str(time.localtime()[3])+str(time.localtime()[4])+str(time.localtime()[5]) sim_state = sim.get_state() control_vector_length=sim.data.ctrl.__len__() #print("control_vector_length: "+str(control_vector_length)) number_of_task_samples=est_activations.shape[0] real_attempt_positions = np.zeros((number_of_task_samples,2)) real_attempt_activations = np.zeros((number_of_task_samples,3)) chassis_pos=np.zeros(number_of_task_samples,) sim.set_state(sim_state) for ii in range(number_of_task_samples): sim.data.ctrl[:] = est_activations[ii,:] sim.step() current_positions_array = sim.data.qpos[-2:] # current_kinematics_array=np.array( # [sim.data.qpos[0], # sim.data.qvel[0], # sim.data.qacc[0], # sim.data.qpos[1], # sim.data.qvel[1], # sim.data.qacc[1]] # ) chassis_pos[ii]=sim.data.get_geom_xpos("Chassis_frame")[0] real_attempt_positions[ii,:] = current_positions_array real_attempt_activations[ii,:] = sim.data.ctrl if Mj_render: viewer.render() real_attempt_kinematics = positions_to_kinematics_fcn( real_attempt_positions[:,0], real_attempt_positions[:,1], timestep = 0.01) return real_attempt_kinematics, real_attempt_activations, chassis_pos
def test_dampingControl(): model = load_model_from_path("robot.xml") sim = MjSim(model) viewer = MjViewer(sim) timestep = generatePatternedTrajectories.TIMESTEP control_freq = 1/timestep total_time = 2 num_cycles = int(total_time * control_freq) plt.ion() LW = 1.0 fig = plt.figure(figsize=(4,15)) axes = [] lines = [] goals = [] for i in range(7): axes.append(fig.add_subplot(7,1,i+1)) lines.append(axes[i].plot([],[],'b-', lw=LW)[0]) goals.append(axes[i].plot([],[],'r-', lw=LW)[0]) axes[i].set_ylim([-1, 1]) axes[i].set_xlim([0,total_time]) axes[i].set_ylabel("Angle{}(rad)".format(i), fontsize=8) axes[i].set_xlabel("Time(s)", fontsize=8) for test in range(5): q_init = bounds.getRandPosInBounds() qd_goal = np.zeros(7) qd_init = np.random.rand(7) for g in range(7): goals[g].set_ydata(np.ones(num_cycles) * qd_goal[g]) goals[g].set_xdata(np.linspace(0,3,num_cycles)) sim.set_state(MjSimState(time=0,qpos=q_init,qvel=qd_init,act=None,udd_state={})) sim.step() sim_time = 0 for i in range(num_cycles): state = sim.get_state() q = state[1] qd = state[2] sim.data.ctrl[:] = controllers.dampingControl(qd=qd) sim.step() viewer.render() if i % 35 == 0: for a in range(7): lines[a].set_xdata(np.append(lines[a].get_xdata(), sim_time)) lines[a].set_ydata(np.append(lines[a].get_ydata(), qd[a])) fig.canvas.draw() fig.canvas.flush_events() sim_time += timestep if bounds.outOfBounds(q): break for i in range(7): lines[i].set_xdata([]) lines[i].set_ydata([]) time.sleep(1)
class MyEnv(gym.Env): metadata = {'render.modes': ['human']} def __init__(self): self.model = load_model_from_path("xmls/Tesrt.xml") self.sim = MjSim(self.model) self.viewer = None self.sim_state = self.sim.get_state() self.bodynames = [ 'torso1', 'head', 'uwaist', 'lwaist', 'butt', 'right_thigh1', 'right_shin1', 'right_foot_cap1', 'right_foot_cap2', 'left_thigh1', 'left_shin1', 'left_foot_cap1', 'left_foot_cap2', 'right_uarm1', 'right_larm', 'right_hand', 'left_uarm1', 'left_larm', 'left_hand' ] ones_act = np.ones(len(self.sim.data.ctrl)) ones_obs = np.ones(self._get_state().shape[0]) self.action_space = spaces.Box(-ones_act, ones_act) self.observation_space = spaces.Box(-ones_obs, ones_obs) def _get_state(self): torso = [] ret = [] for i in range(len(self.bodynames)): vec = self.sim.data.get_geom_xpos(self.bodynames[i]) if i==0: ret = vec torso = vec if i!=0: ret = np.append(ret, vec-torso) return ret def _get_reward(self): return self.sim.data.get_geom_xpos('head')[2] def _step(self, action): for i in range(len(action)): self.sim.data.ctrl[i] = action[i] * 0.5 self.sim.step() self.sim.step() return self._get_state(),self._get_reward(),False,{} def _reset(self): self.sim.set_state(self.sim_state) return self._get_state() def _render(self, mode = 'human', close = False): if self.viewer is None: self.viewer = MjViewer(self.sim) self.viewer.render()
class MujocoSimulation(Simulation): def __init__(self, name, total_time=1000, recording_time=[0, 1000]): super(MujocoSimulation, self).__init__(name, total_time=total_time, recording_time=recording_time) self.model = load_model_from_path('mujoco_xmls/' + name + '.xml') self.sim = MjSim(self.model) self.initial_state = self.sim.get_state() self.input_size = len(self.sim.data.ctrl) self.reset() self.viewer = None def reset(self): super(MujocoSimulation, self).reset() self.sim.set_state(self.initial_state) def run(self, reset=True): if reset: self.reset() self.sim.set_state(self.initial_state) for i in range(self.total_time): self.sim.data.ctrl[:] = self.ctrl_array[i] self.sim.step() self.trajectory.append(self.sim.get_state()) self.alreadyRun = True def get_trajectory(self, all_info=True): if not self.alreadyRun: self.run() if all_info: return self.trajectory.copy() else: return [x.qpos for x in self.trajectory] def get_recording(self, all_info=True): traj = self.get_trajectory(all_info=all_info) return traj[self.recording_time[0]:self.recording_time[1]] def watch(self, repeat_count=4): if self.viewer is None: self.viewer = MjViewer(self.sim) for _ in range(repeat_count): self.sim.set_state(self.initial_state) for i in range(self.total_time): self.sim.data.ctrl[:] = self.ctrl_array[i] self.sim.step() self.viewer.render() self.run( reset=False ) # so that the trajectory will be compatible with what user watches
class gen3_env(object): def __init__(self): #super(lab_env, self).__init__(env) # The real-world simulator self.model = load_model_from_path(xml_path) self.sim = MjSim(self.model) self.viewer = MjViewer(self.sim) def step(self): self.sim.step() self.viewer.render() print(self.sim.data.qpos) # self.sim.data.ctrl[:] = 0.0 for i in range(len(self.sim.data.qvel)): self.sim.data.qvel[i] = 0.1
class Dribble_Env(object): def __init__(self): self.model = load_model_from_path("./xml/world3.xml") self.sim = MjSim(self.model) # self.viewer = MyMjViewer(self.sim) self.viewer = MjViewer(self.sim) self.max_vel = [-1000, 1000] self.x_motor = 0 self.y_motor = 0 def step(self, action): self.x_motor = np.clip(self.x_motor + ((action % 3) - 1) * 100, -500, 500) self.y_motor = np.clip(self.y_motor + ((action // 3) - 1) * 100, -500, 500) self.sim.data.ctrl[0] = self.x_motor self.sim.data.ctrl[1] = self.y_motor # print("---------------------") # print(self.x_motor,self.y_motor,action) # print(self.sim.data.ctrl) self.sim.step() def get_state(self): robot_x, robot_y = self.sim.data.body_xpos[1][0:2] # TODO robot_xv, robot_yv = self.sim.data.qvel[0:2] ball_x, ball_y = self.sim.data.body_xpos[2][0:2] ball_xv, ball_yv = self.sim.data.qvel[2:4] ball_pos_local = -(robot_x - ball_x), -(robot_y - ball_y) # distance = math.sqrt(ball_pos_local[0]**2 + ball_pos_local[1]**2) return [robot_x, robot_y, ball_pos_local[0], ball_pos_local[1], \ robot_xv, robot_yv, ball_x, ball_y,ball_xv,ball_yv] def check_done(self): ball_x, ball_y = self.get_state()[6:8] if ball_x < -80 and -25 < ball_y < 25: return True else: return False def reset(self): self.x_motor = 0 self.y_motor = 0 self.sim.reset() def render(self): self.viewer.render()
def test_cycle_through_orientations(): panda_kinematics = ikpy_panda_kinematics.panda_kinematics() model = load_model_from_path("robot.xml") sim = MjSim(model) viewer = MjViewer(sim) x_target, y_target, z_target = 0.2, 0.0, 0.5 step1 = np.linspace(0,np.pi/2, 100) step2 = np.linspace(np.pi/2, -np.pi/2, 200) step3 = np.linspace(-np.pi/2, 0, 100) sweep = np.concatenate((step1, step2, step3)) roll = 0 pitch = 0 yaw = np.pi ee_goal = [x_target, y_target, z_target, roll, pitch, yaw] qinit = panda_kinematics.inverse_kinematics(translation=ee_goal[0:3], rpy=ee_goal[3:6]) sim.set_state(MjSimState(time=0,qpos=qinit,qvel=np.zeros(7),act=None,udd_state={})) sim.step() qgoal = qinit q = qinit qd = np.zeros(7) count = -1 num_steps_per_change = 4 for i in range(num_steps_per_change * 400): if i % num_steps_per_change == 0: count += 1 ee_goal = [x_target, y_target, z_target, roll, pitch, sweep[count] + yaw] qgoal = panda_kinematics.inverse_kinematics(translation=ee_goal[0:3], rpy=ee_goal[3:6], init_qpos=q) R1 = panda_kinematics.euler_angles_to_rpy_rotation_matrix(rpy=[roll, pitch, sweep[count] + yaw]) R2 = panda_kinematics.euler_angles_to_rpy_rotation_matrix(panda_kinematics.forward_kinematics(qgoal)[3:6]) # R3 = sim.data.body_xmat[sim.model.body_name2id("right_hand")].reshape(3, 3) print("R1:\n", R1) print("R2:\n", R2) # print("R3:\n", R3) print("EE:", np.around(ee_goal, 3)) print("q: ", np.around(qgoal, 3)) state = sim.get_state() q = state[1] qd = state[2] sim.data.ctrl[:] = controllers.PDControl(q=q,qd=qd,qgoal=qgoal) sim.step() viewer.render() time.sleep(1)
def main(): desp = 'Display Robot' parser = argparse.ArgumentParser(description=desp) parser.add_argument('--robot_file', type=str, default='../xml/simrobot/7dof/peg/robot1.xml') args = parser.parse_args() np.set_printoptions(precision=6, suppress=True) print('Displaying robot from:', os.path.abspath(args.robot_file)) model = load_model_from_path(args.robot_file) sim = MjSim(model, nsubsteps=20) sim.reset() sim.step() viewer = MjViewer(sim) while True: viewer.render()
def test(): l, w, h, t, left, m = sample_microwave(False) cab = build_microwave(l, w, h, t, left, set_pose=[0.9, 0.0, -0.15], set_rot=[0, 0, 0, 1]) model = load_model_from_xml(cab.xml) sim = MjSim(model) viewer = MjViewer(sim) t = 0 sim.data.ctrl[0] = -0.2 while t < 5000: sim.step() viewer.render() t += 1
class ArmTaskEnv: """ States: sensors Actions: actuators """ def __init__(self, model_file): # Load model to MuJoCo self.model_file = model_file self.mj_model = load_model_from_path(model_file) self.sim = MjSim(self.mj_model) #self.sim = mujoco.Physics.from_xml_path(model_file) # Parse state and action spaces self.state_dim = len(self.sim.data.sensordata) self.action_dim = len(self.sim.data.ctrl) # Rendering self.viewer = None def reset(self): self.sim.reset() def step(self, action): # Assign action to actuators self.sim.data.ctrl[:] = action.copy() # Simulate step self.sim.step() # Get next state next_state = self.sim.data.sensordata.copy() return next_state def render(self): if self.viewer is None: self.viewer = MjViewer(self.sim) self.viewer.render() @property def time(self): return self.sim.get_state().time
class MjJacoEnv(object): """docstring for MjJacoEnv.""" def __init__(self, vis=False): super(MjJacoEnv, self).__init__() parent_dir_path = str(pathlib.Path(__file__).parent.absolute()) self.fname = parent_dir_path + '/jaco/jaco.xml' self.model = load_model_from_path(self.fname) self.sim = MjSim(self.model) self.viewer = MjViewer(self.sim) self.vis = vis def step(self, action): for i in range(len(action)): self.sim.data.ctrl[i] = action[i] self.sim.forward() self.sim.step() self.viewer.render() if self.vis else None return self.sim.data.qpos
def test(): # import transforms3d as tf3d from mujoco_py import load_model_from_xml, MjSim, MjViewer from mujoco_py.modder import TextureModder for i in range(200): l, w, h, t, left, m = sample_drawers() drawer = build_drawer(l, w, h, t, left) model = load_model_from_xml(drawer.xml) sim = MjSim(model) viewer = MjViewer(sim) modder = TextureModder(sim) for name in sim.model.geom_names: modder.rand_all(name) t = 0 sim.data.ctrl[0] = 0.05 while t < 1000: sim.step() viewer.render() t += 1
def test(): from mujoco_py import load_model_from_xml, MjSim, MjViewer from mujoco_py.modder import TextureModder l, w, h, t, left, m = sample_toaster() cab = build_toaster(l, w, h, t, left) # print(cab.xml) model = load_model_from_xml(cab.xml) sim = MjSim(model) viewer = MjViewer(sim) modder = TextureModder(sim) for name in sim.model.geom_names: modder.rand_all(name) t = 0 sim.data.ctrl[0] = 0.2 while t < 1000: # for name in sim.model.geom_names: # modder.rand_all(name) sim.step() viewer.render() t += 1
def run(): model = load_model_from_path("/Users/zachyamaoka/Documents/de3_group_project/sim/falling_bar.xml") # model = load_model_from_path("/Users/zachyamaoka/Documents/de3_group_project/sim/6_Bar.xml") sim = MjSim(model) viewer = MjViewer(sim) t = 0 vel_data = [] pos_data = [] acc_data = [] time = [] while True: # sim.data.ctrl[0] = math.cos(t / 10.) * 0.01 # sim.data.ctrl[1] = math.sin(t / 10.) * 0.01 time.append(t*0.01) t += 1 # sim.step() viewer.render() if t > 400: break if t > 100 and os.getenv('TESTING') is not None: break
class MyRobotEnv(gym.Env): def __init__(self): self.model = load_model_from_path("my_envs/mujoco/assets/myrobot.xml") self.sim = MjSim(self.model) self.viewer = MjViewer(self.sim) ''' def __init__(self): pass ''' def step(self, action): pass def reset(self): pass def render(self, mode='human', close=False): t = 0 while True: self.viewer.render() t += 1
print("tosser (slide,hinge): ", qpos[:2]) print("object (z,y,pitch): ", qpos[-3:]) model = load_model_from_path("../../xmls/tosser.xml") sim = MjSim(model) viewer = MjViewer(sim) sim_state = sim.get_state() while True: sim.set_state(sim_state) for i in range(1000): state = sim.get_state() # time, qpos, qvel, act, udd_state = state.time, state.qpos, state.qvel, state.act, state.udd_state # print(time, qpos, qvel) print_state(state) if i < 150: sim.data.ctrl[0] = -0.0 sim.data.ctrl[1] = -0.0 else: sim.data.ctrl[0] = -1.0 sim.data.ctrl[1] = -1.0 sim.step() viewer.render() if os.getenv('TESTING') is not None: break
class Environment(): def __init__(self, model_name, goal_space_train, goal_space_test, project_state_to_end_goal, end_goal_thresholds, initial_state_space, subgoal_bounds, project_state_to_subgoal, subgoal_thresholds, max_actions=1200, num_frames_skip=10, show=False): self.name = model_name # Create Mujoco Simulation self.model = load_model_from_path("./mujoco_files/" + model_name) self.sim = MjSim(self.model) # Set dimensions and ranges of states, actions, and goals in order to configure actor/critic networks if model_name == "pendulum.xml": self.state_dim = 2 * len(self.sim.data.qpos) + len( self.sim.data.qvel) else: self.state_dim = len(self.sim.data.qpos) + len( self.sim.data.qvel ) # State will include (i) joint angles and (ii) joint velocities self.action_dim = len( self.sim.model.actuator_ctrlrange) # low-level action dim self.action_bounds = self.sim.model.actuator_ctrlrange[:, 1] # low-level action bounds self.action_offset = np.zeros((len( self.action_bounds))) # Assumes symmetric low-level action ranges self.end_goal_dim = len(goal_space_test) self.subgoal_dim = len(subgoal_bounds) self.subgoal_bounds = subgoal_bounds # Projection functions self.project_state_to_end_goal = project_state_to_end_goal self.project_state_to_subgoal = project_state_to_subgoal # Convert subgoal bounds to symmetric bounds and offset. Need these to properly configure subgoal actor networks self.subgoal_bounds_symmetric = np.zeros((len(self.subgoal_bounds))) self.subgoal_bounds_offset = np.zeros((len(self.subgoal_bounds))) for i in range(len(self.subgoal_bounds)): self.subgoal_bounds_symmetric[i] = (self.subgoal_bounds[i][1] - self.subgoal_bounds[i][0]) / 2 self.subgoal_bounds_offset[i] = self.subgoal_bounds[i][ 1] - self.subgoal_bounds_symmetric[i] # End goal/subgoal thresholds self.end_goal_thresholds = end_goal_thresholds self.subgoal_thresholds = subgoal_thresholds # Set inital state and goal state spaces self.initial_state_space = initial_state_space self.goal_space_train = goal_space_train self.goal_space_test = goal_space_test self.subgoal_colors = [ "Magenta", "Green", "Red", "Blue", "Cyan", "Orange", "Maroon", "Gray", "White", "Black" ] self.max_actions = max_actions # Implement visualization if necessary self.visualize = show # Visualization boolean if self.visualize: self.viewer = MjViewer(self.sim) self.num_frames_skip = num_frames_skip # Get state, which concatenates joint positions and velocities def get_state(self): if self.name == "pendulum.xml": return np.concatenate([ np.cos(self.sim.data.qpos), np.sin(self.sim.data.qpos), self.sim.data.qvel ]) else: return np.concatenate((self.sim.data.qpos, self.sim.data.qvel)) # Reset simulation to state within initial state specified by user def reset_sim(self): # Reset joint positions and velocities for i in range(len(self.sim.data.qpos)): self.sim.data.qpos[i] = np.random.uniform( self.initial_state_space[i][0], self.initial_state_space[i][1]) for i in range(len(self.sim.data.qvel)): self.sim.data.qvel[i] = np.random.uniform( self.initial_state_space[len(self.sim.data.qpos) + i][0], self.initial_state_space[len(self.sim.data.qpos) + i][1]) self.sim.step() # Return state return self.get_state() # Execute low-level action for number of frames specified by num_frames_skip def execute_action(self, action): self.sim.data.ctrl[:] = action for _ in range(self.num_frames_skip): self.sim.step() if self.visualize: self.viewer.render() return self.get_state() # Visualize end goal. This function may need to be adjusted for new environments. def display_end_goal(self, end_goal): # Goal can be visualized by changing the location of the relevant site object. if self.name == "pendulum.xml": self.sim.data.mocap_pos[0] = np.array([ 0.5 * np.sin(end_goal[0]), 0, 0.5 * np.cos(end_goal[0]) + 0.6 ]) elif self.name == "ur5.xml": theta_1 = end_goal[0] theta_2 = end_goal[1] theta_3 = end_goal[2] # shoulder_pos_1 = np.array([0,0,0,1]) upper_arm_pos_2 = np.array([0, 0.13585, 0, 1]) forearm_pos_3 = np.array([0.425, 0, 0, 1]) wrist_1_pos_4 = np.array([0.39225, -0.1197, 0, 1]) # Transformation matrix from shoulder to base reference frame T_1_0 = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.089159], [0, 0, 0, 1]]) # Transformation matrix from upper arm to shoulder reference frame T_2_1 = np.array([[np.cos(theta_1), -np.sin(theta_1), 0, 0], [np.sin(theta_1), np.cos(theta_1), 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) # Transformation matrix from forearm to upper arm reference frame T_3_2 = np.array([[np.cos(theta_2), 0, np.sin(theta_2), 0], [0, 1, 0, 0.13585], [-np.sin(theta_2), 0, np.cos(theta_2), 0], [0, 0, 0, 1]]) # Transformation matrix from wrist 1 to forearm reference frame T_4_3 = np.array([[np.cos(theta_3), 0, np.sin(theta_3), 0.425], [0, 1, 0, 0], [-np.sin(theta_3), 0, np.cos(theta_3), 0], [0, 0, 0, 1]]) # Determine joint position relative to original reference frame # shoulder_pos = T_1_0.dot(shoulder_pos_1) upper_arm_pos = T_1_0.dot(T_2_1).dot(upper_arm_pos_2)[:3] forearm_pos = T_1_0.dot(T_2_1).dot(T_3_2).dot(forearm_pos_3)[:3] wrist_1_pos = T_1_0.dot(T_2_1).dot(T_3_2).dot(T_4_3).dot( wrist_1_pos_4)[:3] joint_pos = [upper_arm_pos, forearm_pos, wrist_1_pos] """ print("\nEnd Goal Joint Pos: ") print("Upper Arm Pos: ", joint_pos[0]) print("Forearm Pos: ", joint_pos[1]) print("Wrist Pos: ", joint_pos[2]) """ for i in range(3): self.sim.data.mocap_pos[i] = joint_pos[i] else: assert False, "Provide display end goal function in environment.py file" # Function returns an end goal def get_next_goal(self, test): end_goal = np.zeros((len(self.goal_space_test))) if self.name == "ur5.xml": goal_possible = False while not goal_possible: end_goal = np.zeros(shape=(self.end_goal_dim, )) end_goal[0] = np.random.uniform(self.goal_space_test[0][0], self.goal_space_test[0][1]) end_goal[1] = np.random.uniform(self.goal_space_test[1][0], self.goal_space_test[1][1]) end_goal[2] = np.random.uniform(self.goal_space_test[2][0], self.goal_space_test[2][1]) # Next need to ensure chosen joint angles result in achievable task (i.e., desired end effector position is above ground) theta_1 = end_goal[0] theta_2 = end_goal[1] theta_3 = end_goal[2] # shoulder_pos_1 = np.array([0,0,0,1]) upper_arm_pos_2 = np.array([0, 0.13585, 0, 1]) forearm_pos_3 = np.array([0.425, 0, 0, 1]) wrist_1_pos_4 = np.array([0.39225, -0.1197, 0, 1]) # Transformation matrix from shoulder to base reference frame T_1_0 = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.089159], [0, 0, 0, 1]]) # Transformation matrix from upper arm to shoulder reference frame T_2_1 = np.array([[np.cos(theta_1), -np.sin(theta_1), 0, 0], [np.sin(theta_1), np.cos(theta_1), 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) # Transformation matrix from forearm to upper arm reference frame T_3_2 = np.array([[np.cos(theta_2), 0, np.sin(theta_2), 0], [0, 1, 0, 0.13585], [-np.sin(theta_2), 0, np.cos(theta_2), 0], [0, 0, 0, 1]]) # Transformation matrix from wrist 1 to forearm reference frame T_4_3 = np.array([[np.cos(theta_3), 0, np.sin(theta_3), 0.425], [0, 1, 0, 0], [-np.sin(theta_3), 0, np.cos(theta_3), 0], [0, 0, 0, 1]]) forearm_pos = T_1_0.dot(T_2_1).dot(T_3_2).dot( forearm_pos_3)[:3] wrist_1_pos = T_1_0.dot(T_2_1).dot(T_3_2).dot(T_4_3).dot( wrist_1_pos_4)[:3] # Make sure wrist 1 pos is above ground so can actually be reached if np.absolute(end_goal[0]) > np.pi / 4 and forearm_pos[ 2] > 0.05 and wrist_1_pos[2] > 0.15: goal_possible = True elif not test and self.goal_space_train is not None: for i in range(len(self.goal_space_train)): end_goal[i] = np.random.uniform(self.goal_space_train[i][0], self.goal_space_train[i][1]) else: assert self.goal_space_test is not None, "Need goal space for testing. Set goal_space_test variable in \"design_env.py\" file" for i in range(len(self.goal_space_test)): end_goal[i] = np.random.uniform(self.goal_space_test[i][0], self.goal_space_test[i][1]) # Visualize End Goal self.display_end_goal(end_goal) return end_goal # Visualize all subgoals def display_subgoals(self, subgoals): # Display up to 10 subgoals and end goal if len(subgoals) <= 11: subgoal_ind = 0 else: subgoal_ind = len(subgoals) - 11 for i in range(1, min(len(subgoals), 11)): if self.name == "pendulum.xml": self.sim.data.mocap_pos[i] = np.array([ 0.5 * np.sin(subgoals[subgoal_ind][0]), 0, 0.5 * np.cos(subgoals[subgoal_ind][0]) + 0.6 ]) # Visualize subgoal self.sim.model.site_rgba[i][3] = 1 subgoal_ind += 1 elif self.name == "ur5.xml": theta_1 = subgoals[subgoal_ind][0] theta_2 = subgoals[subgoal_ind][1] theta_3 = subgoals[subgoal_ind][2] # shoulder_pos_1 = np.array([0,0,0,1]) upper_arm_pos_2 = np.array([0, 0.13585, 0, 1]) forearm_pos_3 = np.array([0.425, 0, 0, 1]) wrist_1_pos_4 = np.array([0.39225, -0.1197, 0, 1]) # Transformation matrix from shoulder to base reference frame T_1_0 = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.089159], [0, 0, 0, 1]]) # Transformation matrix from upper arm to shoulder reference frame T_2_1 = np.array([[np.cos(theta_1), -np.sin(theta_1), 0, 0], [np.sin(theta_1), np.cos(theta_1), 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) # Transformation matrix from forearm to upper arm reference frame T_3_2 = np.array([[np.cos(theta_2), 0, np.sin(theta_2), 0], [0, 1, 0, 0.13585], [-np.sin(theta_2), 0, np.cos(theta_2), 0], [0, 0, 0, 1]]) # Transformation matrix from wrist 1 to forearm reference frame T_4_3 = np.array([[np.cos(theta_3), 0, np.sin(theta_3), 0.425], [0, 1, 0, 0], [-np.sin(theta_3), 0, np.cos(theta_3), 0], [0, 0, 0, 1]]) # Determine joint position relative to original reference frame # shoulder_pos = T_1_0.dot(shoulder_pos_1) upper_arm_pos = T_1_0.dot(T_2_1).dot(upper_arm_pos_2)[:3] forearm_pos = T_1_0.dot(T_2_1).dot(T_3_2).dot( forearm_pos_3)[:3] wrist_1_pos = T_1_0.dot(T_2_1).dot(T_3_2).dot(T_4_3).dot( wrist_1_pos_4)[:3] joint_pos = [upper_arm_pos, forearm_pos, wrist_1_pos] """ print("\nSubgoal %d Joint Pos: " % i) print("Upper Arm Pos: ", joint_pos[0]) print("Forearm Pos: ", joint_pos[1]) print("Wrist Pos: ", joint_pos[2]) """ # Designate site position for upper arm, forearm and wrist for j in range(3): self.sim.data.mocap_pos[3 + 3 * (i - 1) + j] = np.copy( joint_pos[j]) self.sim.model.site_rgba[3 + 3 * (i - 1) + j][3] = 1 # print("\nLayer %d Predicted Pos: " % i, wrist_1_pos[:3]) subgoal_ind += 1 else: # Visualize desired gripper position, which is elements 18-21 in subgoal vector self.sim.data.mocap_pos[i] = subgoals[subgoal_ind] # Visualize subgoal self.sim.model.site_rgba[i][3] = 1 subgoal_ind += 1
def main(data_dir): model = load_model_from_path("robot.xml") sim = MjSim(model) viewer = MjViewer(sim) initial_state = MjSimState(time=0, qpos=np.array([ 0, -np.pi / 4, 0, -3 * np.pi / 4 + 1, 0, np.pi / 2, np.pi / 4 ]), qvel=np.zeros(7), act=None, udd_state={}) sim.set_state(initial_state) sim.step() traj_count = 0 control_state = None while (traj_count < NUM_TRAJECTORIES and control_state != FINISHED): control_state = ACCELERATING outputFile = None initial_q = sim.get_state()[q_INDEX] velGen = randVelGen.RandVelGenerator() qd_des = velGen.generatePatternVel() coming_back_time = 0.0 time = 0 while control_state != FINISHED: state = sim.get_state() q = state[q_INDEX] qd = state[qd_INDEX] boundViolations = bounds.getBoundViolations(q) # RD = TB = bounds.tableBoundViolation(sim) OB = bounds.outOfBounds(q) DA = helperFun.moving(qd) DC = helperFun.stopped(qd) DD = DC DB = coming_back_time > RETURN_TIME FN = traj_count >= NUM_TRAJECTORIES prev_state = control_state # transition block if control_state == ACCELERATING: if not TB and not OB and not DA: control_state = ACCELERATING elif TB: control_state = COMING_BACK_IN_BOUNDS coming_back_time = 0.0 elif not TB and OB: control_state = DAMPING curBoundViolations = bounds.getBoundViolations(q) velGen.setJointDirections(curBoundViolations) elif not TB and not OB and DA: control_state = COASTING traj_count += 1 outputFile = open(data_dir + helperFun.getUniqueFileName("traj"), mode='x') outputWriter = csv.writer(outputFile, delimiter=',') print_count(traj_count) else: control_state = FINISHED print("Unknown transistion! ACCELERATING") elif control_state == COASTING: if not FN and not TB and not OB and DC: control_state = ACCELERATING qd_des = velGen.generatePatternVel() outputFile.close() elif not FN and TB: control_state = COMING_BACK_IN_BOUNDS coming_back_time = 0 outputFile.close() elif not FN and not TB and OB: control_state = DAMPING outputFile.close() curBoundViolations = bounds.getBoundViolations(q) velGen.setJointDirections(curBoundViolations) elif FN: control_state = FINISHED outputFile.close() elif not FN and not TB and not OB and not DC: control_state = COASTING else: control_state = FINISHED print("Unknown transition! COASTING") outputFile.close() elif control_state == DAMPING: if not TB and not DD: control_state = DAMPING elif TB: control_state = COMING_BACK_IN_BOUNDS coming_back_time = 0.0 elif not TB and DD: control_state = ACCELERATING qd_des = velGen.generatePatternVel() else: control_state = FINISHED print("Unknow transition! DAMPING") elif control_state == COMING_BACK_IN_BOUNDS: if not DB: control_state = COMING_BACK_IN_BOUNDS elif DB and OB: control_state = DAMPING curBoundViolations = bounds.getBoundViolations(q) velGen.setJointDirections(curBoundViolations) elif DB and not OB: control_state = ACCELERATING qd_des = velGen.generatePatternVel() else: control_state = FINISHED print("Unknown transition! COMING_BACK_IN_BOUNDS") elif control_state == FINISHED: control_state = FINISHED else: control_state = FINISHED print("Got to an invalid state!") # debug states if prev_state != control_state: if control_state == ACCELERATING: print("ACCELERATING") elif control_state == COASTING: print("COASTING") elif control_state == DAMPING: print("DAMPING") elif control_state == COMING_BACK_IN_BOUNDS: print("COMING_BACK_IN_BOUNDS") elif control_state == "FINISHED": print("FINISHED") else: print("In a bad state!") torques = np.zeros(7) if control_state == ACCELERATING: torques = controllers.basicVelControl(qd_des=qd_des, qd_cur=qd) elif control_state == COASTING: data = np.concatenate( (q, qd, data_calc.get_3D_data(sim), [time])) outputWriter.writerow(data) torques = controllers.basicVelControl(qd_des=qd_des, qd_cur=qd) elif control_state == DAMPING: torques = controllers.dampingControl(qd) elif control_state == COMING_BACK_IN_BOUNDS: coming_back_time += TIMESTEP torques = controllers.moveAwayFromTable(q=q, qd=qd) elif control_state == FINISHED: outputFile.close() break else: print("Got to an invalid state!") control_state = FINISHED break sim.data.ctrl[:] = torques sim.step() viewer.render() time += TIMESTEP
class GripperTester: """ A class that is used to test gripper Args: gripper (GripperModel): A gripper instance to be tested pos (str): (x y z) position to place the gripper in string form, e.g. '0 0 0.3' quat (str): rotation to apply to gripper in string form, e.g. '0 0 1 0' to flip z axis gripper_low_pos (float): controls the gipper y position, larger -> higher gripper_high_pos (float): controls the gipper y high position larger -> higher, must be larger than gripper_low_pos box_size (None or 3-tuple of int): the size of the box to grasp, None defaults to [0.02, 0.02, 0.02] box_density (int): the density of the box to grasp step_time (int): the interval between two gripper actions render (bool): if True, show rendering """ def __init__( self, gripper, pos, quat, gripper_low_pos, gripper_high_pos, box_size=None, box_density=10000, step_time=400, render=True, ): # define viewer self.viewer = None world = MujocoWorldBase() # Add a table arena = TableArena(table_full_size=(0.4, 0.4, 0.1), table_offset=(0, 0, 0.1), has_legs=False) world.merge(arena) # Add a gripper self.gripper = gripper # Create another body with a slider joint to which we'll add this gripper gripper_body = ET.Element("body") gripper_body.set("pos", pos) gripper_body.set("quat", quat) # flip z gripper_body.append(new_joint(name="gripper_z_joint", type="slide", axis="0 0 -1", damping="50")) # Add all gripper bodies to this higher level body for body in gripper.worldbody: gripper_body.append(body) # Merge the all of the gripper tags except its bodies world.merge(gripper, merge_body=None) # Manually add the higher level body we created world.worldbody.append(gripper_body) # Create a new actuator to control our slider joint world.actuator.append(new_actuator(joint="gripper_z_joint", act_type="position", name="gripper_z", kp="500")) # Add an object for grasping # density is in units kg / m3 TABLE_TOP = [0, 0, 0.09] if box_size is None: box_size = [0.02, 0.02, 0.02] box_size = np.array(box_size) self.cube = BoxObject( name="object", size=box_size, rgba=[1, 0, 0, 1], friction=[1, 0.005, 0.0001], density=box_density ) object_pos = np.array(TABLE_TOP + box_size * [0, 0, 1]) mujoco_object = self.cube.get_obj() # Set the position of this object mujoco_object.set("pos", array_to_string(object_pos)) # Add our object to the world body world.worldbody.append(mujoco_object) # add reference objects for x and y axes x_ref = BoxObject( name="x_ref", size=[0.01, 0.01, 0.01], rgba=[0, 1, 0, 1], obj_type="visual", joints=None ).get_obj() x_ref.set("pos", "0.2 0 0.105") world.worldbody.append(x_ref) y_ref = BoxObject( name="y_ref", size=[0.01, 0.01, 0.01], rgba=[0, 0, 1, 1], obj_type="visual", joints=None ).get_obj() y_ref.set("pos", "0 0.2 0.105") world.worldbody.append(y_ref) self.world = world self.render = render self.simulation_ready = False self.step_time = step_time self.cur_step = 0 if gripper_low_pos > gripper_high_pos: raise ValueError( "gripper_low_pos {} is larger " "than gripper_high_pos {}".format(gripper_low_pos, gripper_high_pos) ) self.gripper_low_pos = gripper_low_pos self.gripper_high_pos = gripper_high_pos def start_simulation(self): """ Starts simulation of the test world """ model = self.world.get_model(mode="mujoco_py") self.sim = MjSim(model) if self.render: self.viewer = MjViewer(self.sim) self.sim_state = self.sim.get_state() # For gravity correction gravity_corrected = ["gripper_z_joint"] self._gravity_corrected_qvels = [self.sim.model.get_joint_qvel_addr(x) for x in gravity_corrected] self.gripper_z_id = self.sim.model.actuator_name2id("gripper_z") self.gripper_z_is_low = False self.gripper_actuator_ids = [self.sim.model.actuator_name2id(x) for x in self.gripper.actuators] self.gripper_is_closed = True self.object_id = self.sim.model.body_name2id(self.cube.root_body) object_default_pos = self.sim.data.body_xpos[self.object_id] self.object_default_pos = np.array(object_default_pos, copy=True) self.reset() self.simulation_ready = True def reset(self): """ Resets the simulation to the initial state """ self.sim.set_state(self.sim_state) self.cur_step = 0 def close(self): """ Close the viewer if it exists """ if self.viewer is not None: self.viewer.close() def step(self): """ Forward the simulation by one timestep Raises: RuntimeError: if start_simulation is not yet called. """ if not self.simulation_ready: raise RuntimeError("Call start_simulation before calling step") if self.gripper_z_is_low: self.sim.data.ctrl[self.gripper_z_id] = self.gripper_low_pos else: self.sim.data.ctrl[self.gripper_z_id] = self.gripper_high_pos if self.gripper_is_closed: self._apply_gripper_action(1) else: self._apply_gripper_action(-1) self._apply_gravity_compensation() self.sim.step() if self.render: self.viewer.render() self.cur_step += 1 def _apply_gripper_action(self, action): """ Applies binary gripper action Args: action (int): Action to apply. Should be -1 (open) or 1 (closed) """ gripper_action_actual = self.gripper.format_action(np.array([action])) # rescale normalized gripper action to control ranges ctrl_range = self.sim.model.actuator_ctrlrange[self.gripper_actuator_ids] bias = 0.5 * (ctrl_range[:, 1] + ctrl_range[:, 0]) weight = 0.5 * (ctrl_range[:, 1] - ctrl_range[:, 0]) applied_gripper_action = bias + weight * gripper_action_actual self.sim.data.ctrl[self.gripper_actuator_ids] = applied_gripper_action def _apply_gravity_compensation(self): """ Applies gravity compensation to the simulation """ self.sim.data.qfrc_applied[self._gravity_corrected_qvels] = self.sim.data.qfrc_bias[ self._gravity_corrected_qvels ] def loop(self, total_iters=1, test_y=False, y_baseline=0.01): """ Performs lower, grip, raise and release actions of a gripper, each separated with T timesteps Args: total_iters (int): Iterations to perform before exiting test_y (bool): test if object is lifted y_baseline (float): threshold for determining that object is lifted """ seq = [(False, False), (True, False), (True, True), (False, True)] for cur_iter in range(total_iters): for cur_plan in seq: self.gripper_z_is_low, self.gripper_is_closed = cur_plan for step in range(self.step_time): self.step() if test_y: if not self.object_height > y_baseline: raise ValueError( "object is lifed by {}, ".format(self.object_height) + "not reaching the requirement {}".format(y_baseline) ) @property def object_height(self): """ Queries the height (z) of the object compared to on the ground Returns: float: Object height relative to default (ground) object position """ return self.sim.data.body_xpos[self.object_id][2] - self.object_default_pos[2]
class WindySlope(gym.Env): def __init__(self, model, mode, hertz=25, should_render=True, should_screenshot=False): self.hertz = hertz self.steps = 0 self.should_render = should_render self.should_screenshot = should_screenshot self.nsubsteps = int(MAX_TIME / model.opt.timestep / 100) self.viewer = None self.model = model self.mode = mode self.enabled = True self.metadata = {'render.modes': 'rgb_array'} self.should_record = False def close(self): pass @property def observation_space(self): return Box(low=-np.inf, high=np.inf, shape=(18, )) @property def action_space(self): return Box(low=-np.inf, high=np.inf, shape=(0, )) @property def model(self): return self._model @model.setter def model(self, model): self._model = model self.sim = MjSim(model) self.data = self.sim.data if self.should_render: if self.viewer: self.viewer.sim = sim return self.viewer = MjViewer(self.sim) self.viewer.cam.azimuth = 45 self.viewer.cam.elevation = -20 self.viewer.cam.distance = 25 self.viewer.cam.lookat[:] = [0, 0, -2] self.viewer.scn.flags[3] = 0 def reset(self): self.sim.reset() self.steps = 0 self.sim.forward() obs = self.get_observations(self.model, self.data) return obs def get_observations(self, model, data): self.sim.forward() obs = [] name = 'box' pos = data.get_body_xpos(name) xmat = data.get_body_xmat(name).reshape(-1) velp = data.get_body_xvelp(name) velr = data.get_body_xvelr(name) for x in [pos, xmat, velp, velr]: obs.extend(x.copy()) obs = np.array(obs, dtype=np.float32) return obs def screenshot(self, image_path): self.viewer.hide_overlay = True self.viewer.render() width, height = 2560, 1440 #width, height = 1,1 img = self.viewer.read_pixels(width, height, depth=False) # original image is upside-down, so flip it img = img[::-1, :, :] imageio.imwrite(image_path, img) def step(self, action): nsubsteps = self.nsubsteps for _ in range(nsubsteps): self.sim.step() self.render() self.steps += 1 return self.get_observations(self.model, self.data), 1, self.steps == 100, {} def render(self, mode=None): if self.should_render: self.viewer._overlay.clear() def nothing(): return self.viewer._create_full_overlay = nothing wind = model.opt.wind[0] self.viewer.add_overlay(1, "Wind", "{:.2f}".format(wind)) self.viewer.render() if self.should_record: width, height = 2560, 1440 img = self.viewer.read_pixels(width, height, depth=False) # original image is upside-down, so flip it img = img[::-1, :, :] return img def euler2quat(self, euler): euler = np.asarray(euler, dtype=np.float64) assert euler.shape[-1] == 3, "Invalid shape euler {}".format(euler) ai, aj, ak = euler[..., 2] / 2, -euler[..., 1] / 2, euler[..., 0] / 2 si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak) ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak) cc, cs = ci * ck, ci * sk sc, ss = si * ck, si * sk quat = np.empty(euler.shape[:-1] + (4, ), dtype=np.float64) quat[..., 0] = cj * cc + sj * ss quat[..., 3] = cj * sc - sj * cs quat[..., 2] = -(cj * ss + sj * cc) quat[..., 1] = cj * cs - sj * sc return quat def degrees2radians(self, degrees): return degrees * np.pi / 180
class my_env(parameterized.TestCase): def __init__(self): # super(lab_env, self).__init__(env) # 导入xml文档 self.model = load_model_from_path("assets/simpleEE_4box.xml") # 调用MjSim构建一个basic simulation self.sim = MjSim(model=self.model) self.sim = MjSim(self.model) self.viewer = MjViewer(self.sim) self.viewer._run_speed = 0.001 self.timestep = 0 # Sawyer Peg #self.init_qpos = np.array([-0.305, -0.83, 0.06086, 1.70464, -0.02976, 0.62496, -0.04712]) # Flexiv Peg self.init_qpos = np.array([-0.22, -0.43, 0.449, -2, -0.25, 0.799, 0.99]) for i in range(len(self.sim.data.qpos)): self.sim.data.qpos[i] = self.init_qpos[i] self.testQposFromSitePose( (np.array([0.57, 0.075, 0.08]), np.array([0.000000e+00, 1.000000e+00, 0.000000e+00, 6.123234e-17])), _INPLACE, True) print(self.sim.data.ctrl) print(self.sim.data.qpos) def get_state(self): self.sim.get_state() # 如果定义了相机 # self.sim.data.get_camera_xpos('[camera name]') def reset(self): self.sim.reset() self.timestep = 0 def step(self): # self.testQposFromSitePose((np.array([0.605, 0.075, 0.03]), np.array([0.000000e+00, 1.000000e+00, 0.000000e+00, 6.123234e-17])), _INPLACE) x=random.uniform(0.415, 0.635) y=random.uniform(-0.105, 0.105) self.testQposFromSitePose( (np.array([x, y, 0.045]), np.array([0.000000e+00, 1.000000e+00, 0.000000e+00, 6.123234e-17])), _INPLACE) # self.testQposFromSitePose( # (None, np.array([0.000000e+00, 1.000000e+00, 0.000000e+00, 6.123234e-17])), # _INPLACE, True) self.sim.step() # self.sim.data.ctrl[0] += 0.01 # print(self.sim.data.ctrl) # pdb.set_trace() # print(self.sim.data.qpos) print("sensordata", self.sim.data.sensordata) # self.viewer.add_overlay(const.GRID_TOPRIGHT, " ", SESSION_NAME) self.viewer.render() self.timestep += 1 def create_viewer(self, run_speed=0.0005): self.viewer = MjViewer(self.sim) self.viewer._run_speed = run_speed # self.viewer._hide_overlay = HIDE_OVERLAY # self.viewer.vopt.frame = DISPLAY_FRAME # self.viewer.cam.azimuth = CAM_AZIMUTH # self.viewer.cam.distance = CAM_DISTANCE # self.viewer.cam.elevation = CAM_ELEVATION def testQposFromSitePose(self, target, inplace, qpos_flag=False): physics = mujoco.Physics.from_xml_string(FlexivPeg_XML) target_pos, target_quat = target count = 0 physics2 = physics.copy(share_model=True) resetter = _ResetArm(seed=0) while True: result = ik.qpos_from_site_pose( physics=physics2, site_name=_SITE_NAME, target_pos=target_pos, target_quat=target_quat, joint_names=_JOINTS, tol=_TOL, max_steps=_MAX_STEPS, inplace=inplace, ) if result.success: break elif count < _MAX_RESETS: resetter(physics2) count += 1 else: raise RuntimeError( 'Failed to find a solution within %i attempts.' % _MAX_RESETS) self.assertLessEqual(result.steps, _MAX_STEPS) self.assertLessEqual(result.err_norm, _TOL) # pdb.set_trace() physics.data.qpos[:] = result.qpos for i in range(len(self.sim.data.qpos)): if qpos_flag: self.sim.data.qpos[i]=physics.data.qpos[i] else: self.sim.data.ctrl[i] = physics.data.qpos[i] # print(physics.data.qpos) mjlib.mj_fwdPosition(physics.model.ptr, physics.data.ptr) if target_pos is not None: pos = physics.named.data.site_xpos[_SITE_NAME] np.testing.assert_array_almost_equal(pos, target_pos) if target_quat is not None: xmat = physics.named.data.site_xmat[_SITE_NAME] quat = np.empty_like(target_quat) mjlib.mju_mat2Quat(quat, xmat) quat /= quat.ptp() # Normalize xquat so that its max-min range is 1
class MujocoEnv(gym.Env): def __init__(self, model_path, frame_skip=1, action_noise=0.0, random_init_state=True): if model_path.startswith("/"): fullpath = model_path else: fullpath = os.path.join(os.path.dirname(__file__), "assets", model_path) if not path.exists(fullpath): raise IOError("File %s does not exist" % fullpath) self.frame_skip = frame_skip self.model = load_model_from_path(fullpath) self.sim = MjSim(self.model) self.data = self.sim.data self.viewer = None self.init_qpos = self.data.qpos.ravel().copy() self.init_qvel = self.data.qvel.ravel().copy() self.init_qacc = self.data.qacc.ravel().copy() self.init_ctrl = self.data.ctrl.ravel().copy() self.qpos_dim = self.init_qpos.size self.qvel_dim = self.init_qvel.size self.ctrl_dim = self.init_ctrl.size self.action_noise = action_noise self.random_init_state = random_init_state """ if "init_qpos" in self.model.numeric_names: init_qpos_id = self.model.numeric_names.index("init_qpos") addr = self.model.numeric_adr.flat[init_qpos_id] size = self.model.numeric_size.flat[init_qpos_id] init_qpos = self.model.numeric_data.flat[addr:addr + size] self.init_qpos = init_qpos """ self.dcom = None self.current_com = None self.reset() super(MujocoEnv, self).__init__() @property def action_space(self): bounds = self.model.actuator_ctrlrange.copy() lb = bounds[:, 0] ub = bounds[:, 1] return spaces.Box(lb, ub) @property def observation_space(self): shp = self.get_current_obs().shape ub = BIG * np.ones(shp) return spaces.Box(ub * -1, ub) @property def action_bounds(self): return self.action_space.low, self.action_space.high def reset_mujoco(self, init_state=None): if init_state is None: if self.random_init_state: qp = self.init_qpos.copy() + \ np.random.normal(size=self.init_qpos.shape) * 0.01 qv = self.init_qvel.copy() + \ np.random.normal(size=self.init_qvel.shape) * 0.1 else: qp = self.init_qpos.copy() qv = self.init_qvel.copy() qacc = self.init_qacc.copy() ctrl = self.init_ctrl.copy() else: pass """ start = 0 for datum_name in ["qpos", "qvel", "qacc", "ctrl"]: datum = getattr(self.data, datum_name) datum_dim = datum.shape[0] datum = init_state[start: start + datum_dim] setattr(self.model.data, datum_name, datum) start += datum_dim """ self.set_state(qp, qv) def reset(self, init_state=None): # self.reset_mujoco(init_state) self.sim.reset() self.sim.forward() self.current_com = self.data.subtree_com[0] self.dcom = np.zeros_like(self.current_com) return self.get_current_obs() def set_state(self, qpos, qvel, qacc): assert qpos.shape == (self.qpos_dim, ) and qvel.shape == ( self.qvel_dim, ) and qacc.shape == (self.qacc_dim, ) state = self.sim.get_state() for i in range(self.model.nq): state.qpos[i] = qpos[i] for i in range(self.model.nv): state.qvel[i] = qvel[i] self.sim.set_state(state) self.sim.forward() def get_current_obs(self): return self._get_full_obs() def _get_full_obs(self): data = self.data cdists = np.copy(self.model.geom_margin).flat for c in self.model.data.contact: cdists[c.geom2] = min(cdists[c.geom2], c.dist) obs = np.concatenate([ data.qpos.flat, data.qvel.flat, # data.cdof.flat, data.cinert.flat, data.cvel.flat, # data.cacc.flat, data.qfrc_actuator.flat, data.cfrc_ext.flat, data.qfrc_constraint.flat, cdists, # data.qfrc_bias.flat, # data.qfrc_passive.flat, self.dcom.flat, ]) return obs @property def _state(self): return np.concatenate([self.data.qpos.flat, self.data.qvel.flat]) @property def _full_state(self): return np.concatenate([ self.data.qpos, self.data.qvel, self.data.qacc, self.data.ctrl, ]).ravel() def inject_action_noise(self, action): # generate action noise noise = self.action_noise * \ np.random.normal(size=action.shape) # rescale the noise to make it proportional to the action bounds lb, ub = self.action_bounds noise = 0.5 * (ub - lb) * noise return action + noise def forward_dynamics(self, action): ctrl = self.inject_action_noise(action) for i in range(self.model.nu): self.sim.data.ctrl[i] = ctrl[i] for _ in range(self.frame_skip): self.sim.step() new_com = self.data.subtree_com[0] self.dcom = new_com - self.current_com self.current_com = new_com def get_viewer(self, config=None): if self.viewer is None: self.viewer = MjViewer(self.sim) # self.viewer.start() # self.viewer.set_model(self.model) if config is not None: pass # self.viewer.set_window_pose(config["xpos"], config["ypos"]) # self.viewer.set_window_size(config["width"], config["height"]) # self.viewer.set_window_title(config["title"]) return self.viewer def render(self, close=False, mode='human', config=None): if mode == 'human': # viewer = self.get_viewer(config=config) try: self.viewer.render() except: self.get_viewer(config=config) self.viewer.render() elif mode == 'rgb_array': viewer = self.get_viewer(config=config) viewer.loop_once() # self.get_viewer(config=config).render() data, width, height = self.get_viewer(config=config).get_image() return np.fromstring(data, dtype='uint8').reshape(height, width, 3)[::-1, :, :] if close: self.stop_viewer() # def start_viewer(self): # viewer = self.get_viewer() # if not viewer.running: # viewer.start() # # def stop_viewer(self): # if self.viewer: # self.viewer.finish() # self.viewer = None # def release(self): # # temporarily alleviate the issue (but still some leak) # from learning_to_adapt.mujoco_py.mjlib import mjlib # mjlib.mj_deleteModel(self.model._wrapped) # mjlib.mj_deleteData(self.data._wrapped) def get_body_xmat(self, body_name): idx = self.model.body_names.index(body_name) return self.data.ximat[idx].reshape((3, 3)) def get_body_com(self, body_name): return self.data.get_body_xpos(body_name) def get_body_comvel(self, body_name): idx = self.model.body_names.index(body_name) ## _compute_subtree body_vels = np.zeros((self.model.nbody, 6)) # bodywise quantities mass = self.model.body_mass.flatten() for i in range(self.model.nbody): # body velocity # Compute object 6D velocity in object-centered frame, world/local orientation. # mj_objectVelocity(const mjModel* m, const mjData* d, int objtype, int objid, mjtMum* res, int flg_local) mujoco_py.cymj._mj_objectVelocity(self.model, self.data, 1, i, body_vels[i], 0) lin_moms = body_vels[:, 3:] * mass.reshape((-1, 1)) # init subtree mass body_parentid = self.model.body_parentid # subtree com and com_vel for i in range(self.model.nbody - 1, -1, -1): if i > 0: parent = body_parentid[i] # add scaled velocities lin_moms[parent] += lin_moms[i] # accumulate mass mass[parent] += mass[i] return_ = lin_moms / mass.reshape((-1, 1)) return return_[idx] # return self.model.body_comvels[idx] # def get_body_comvel(self, body_name): # idx = self.model.body_names.index(body_name) # # return self.model.body_comvels[idx] # def print_stats(self): # super(MujocoEnv, self).print_stats() # print("qpos dim:\t%d" % len(self.data.qpos)) def action_from_key(self, key): raise NotImplementedError # def set_state_tmp(self, state, restore=True): # if restore: # prev_pos = self.data.qpos # prev_qvel = self.data.qvel # prev_ctrl = self.data.ctrl # prev_act = self.data.act # qpos, qvel = self.decode_state(state) # self.model.data.qpos = qpos # self.model.data.qvel = qvel # self.model.forward() # yield # if restore: # self.data.qpos = prev_pos # self.data.qvel = prev_qvel # self.data.ctrl = prev_ctrl # self.data.act = prev_act # self.model.forward() def get_param_values(self): return {} def set_param_values(self, values): pass
#!/usr/bin/env python3 """ Shows how to use render callback. """ from mujoco_py import load_model_from_path, MjSim, MjViewer from mujoco_py.modder import TextureModder import os modder = None def render_callback(sim, viewer): global modder if modder is None: modder = TextureModder(sim) for name in sim.model.geom_names: modder.rand_all(name) model = load_model_from_path("xmls/fetch/main.xml") sim = MjSim(model, render_callback=render_callback) viewer = MjViewer(sim) t = 0 while True: viewer.render() t += 1 if t > 100 and os.getenv('TESTING') is not None: break