def test_tableBoundViolation(): model = load_model_from_path("robot.xml") sim = MjSim(model) test1 = np.array([-0.0031, -0.9718, -0.7269, -2.4357, -0.0157, 1.5763, 0.7303]) #False test2 = np.array([0.0264, -0.0772, 0.1924, -2.8478, -0.0273, 2.8339, 0.7566]) #True test3 = np.array([-1.4870, -1.7289, 1.6138, -1.9814, -0.9856, 1.9304, 0.9981]) #True test4 = np.array([-0.5250, -0.6410, 0.1893, -1.3827, -0.2573, 2.1356, 0.7116]) #False test5 = np.array([-0.0133, 0.9846, 0.0365, -1.5491, 2.8629, 0.7630, 0.6254]) #True qd = np.zeros(7) state = MjSimState(time=0,qpos=test1,qvel=qd,act=None,udd_state={}) sim.set_state(state) sim.step() print("Test1:", "Passed" if not bounds.tableBoundViolation(sim) else "Failed") state = MjSimState(time=0,qpos=test2,qvel=qd,act=None,udd_state={}) sim.set_state(state) sim.step() print("Test2:", "Passed" if bounds.tableBoundViolation(sim) else "Failed") state = MjSimState(time=0,qpos=test3,qvel=qd,act=None,udd_state={}) sim.set_state(state) sim.step() print("Test3:", "Passed" if bounds.tableBoundViolation(sim) else "Failed") state = MjSimState(time=0,qpos=test4,qvel=qd,act=None,udd_state={}) sim.set_state(state) sim.step() print("Test4:", "Passed" if not bounds.tableBoundViolation(sim) else "Failed") state = MjSimState(time=0,qpos=test5,qvel=qd,act=None,udd_state={}) sim.set_state(state) sim.step() print("Test5:", "Passed" if bounds.tableBoundViolation(sim) else "Failed")
def rollout_from_state(opt, env): t = 'textured' if opt.textured or opt.collision else 'normal' os.makedirs(opt.data_root + '/' + opt.dataset_name + '/state_rollout_{}'.format(t), exist_ok=True) save_path = opt.data_root + '/' + opt.dataset_name + '/state_rollout_{}'.format(t) states = np.load(opt.states_file + '.npy') if opt.textured: with open('mujoco/model.xml', 'r') as f: xml = f.read() d = os.path.dirname(os.path.abspath(__file__)) env.unwrapped.reset_from_xml_string(xml.format(base=d)) else: env.reset() fig = plt.figure() ax = fig.add_subplot(1,1,1) for state_index in range(states.shape[0]): if opt.noisy: states[state_index] += np.random.randn(states[state_index].shape) state = MjSimState.from_flattened(states[state_index], env.unwrapped.sim) env.unwrapped.sim.set_state(state) env.unwrapped.sim.forward() obs = env.unwrapped._get_observation()['image'][::-1] obs = Image.fromarray(obs).convert('RGB') obs.save(save_path + '/from_states_{}.jpg'.format(state_index)) """
def set_state(self, qpos, qvel): assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,) old_state = self.sim.get_state() new_state = MjSimState(old_state.time, qpos, qvel, old_state.act, old_state.udd_state) self.sim.set_state(new_state) self.sim.forward()
def frames(self): if self._frames is None: sim_params = GTSimParams(env_name='ObstaclePush-v0', num_action_repeat=5, output_resolution=128, render_goal=True) sim = GTSimulator(sim_params) init_state = self.data['init_state'] # edit MuJoCo environment init state for backward compatibility if len(init_state.qpos) == 12: from mujoco_py import MjSimState new_qpos = np.concatenate([init_state.qpos[:9], np.array([10, 10, 0]), init_state.qpos[-3:]]) new_qvel = np.concatenate([init_state.qvel[:9], np.array([0, 0, 0]), init_state.qvel[-3:]]) init_state = MjSimState(time=0, qpos=new_qpos, qvel=new_qvel, act=init_state.act, udd_state=init_state.udd_state) output_dict = sim.rollout(init_state, np.expand_dims(self.actions, axis=0)) self._frames = output_dict.pred_frames[0] return self._frames
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 _initialize_mujoco_state(self): init_position = np.array(self.init_mean) + np.random.uniform( 0., self.init_noise, 2) init_state = MjSimState(time=0., qpos=init_position, qvel=np.array([0., 0.]), act=None, udd_state={}) self.sim.set_state(init_state)
def reset(self) -> None: """Resets the environment (including the agent) to the initial conditions. """ self.sim.reset() # Set initial position and velocity self.qpos = self.sim.data.qpos.copy() self.qpos[:self.robot_init_qpos.shape[0]] = self.robot_init_qpos qvel = np.zeros(self.sim.data.qvel.shape) mjSimState = MjSimState(time=0.0, qpos=self.qpos, qvel=qvel, act=None, udd_state={}) self.sim.set_state(mjSimState) self.sim.forward()
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)
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 set_initial_robot_state(sim, robot): ''' Used for setting initial state when simulating with mujoco-py directly ''' old_state = sim.get_state() old_qvel = old_state.qvel[robot.dof:] old_qpos = old_state.qpos[ robot.dof:] if robot.gripper.dof < 1 else old_state.qpos[robot.dof - 1:] init_qvel = [0 for _ in range(robot.dof)] init_model_qvel = np.concatenate((init_qvel, old_qvel), axis=0) init_model_qpos = np.concatenate((robot.init_qpos, old_qpos), axis=0) new_state = MjSimState(old_state.time, init_model_qpos, init_model_qvel, old_state.act, old_state.udd_state) sim.set_state(new_state) sim.forward()
def reset(self): """Resets the environment (including the agent) to the initial conditions. """ self.sim.reset() # Set initial position and velocity self.qpos = self.sim.data.qpos.copy() self.qpos[:self.robot_init_qpos.shape[0]] = self.robot_init_qpos if self.random_env: self.callback_randomize_env() qvel = np.zeros(self.sim.data.qvel.shape) mjSimState = MjSimState(time=0.0, qpos=self.qpos, qvel=qvel, act=None, udd_state={}) self.sim.set_state(mjSimState) self.sim.forward() self.terminated = False self.env_step_counter = 0 return self.get_observations()
def set_q(self, joint_pos): """ Sets the value of the robot joints. Args: joint_pos: Value for the robot joints. Returns: No return value """ qpos = self.sim.data.qpos.copy() qpos[:joint_pos.shape[0]] = joint_pos self.qpos = qpos self.qvel = self.sim.data.qvel.copy() self.qvel[:joint_pos.shape[0]] = np.zeros(joint_pos.shape[0]) # self.qvel = np.zeros(self.sim.data.qvel.shape) mjSimState = MjSimState(time=0.0, qpos=self.qpos, qvel=self.qvel, act=None, udd_state={}) self.sim.set_state(mjSimState)
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
def test_sim_state(): model = load_model_from_xml(BASIC_MODEL_XML) foo = 10 d = {"foo": foo, "foo_array": np.array([foo, foo, foo]), "foo_2darray": np.reshape(np.array([foo, foo, foo, foo]), (2, 2)), } def udd_callback(sim): return d sim = MjSim(model, nsubsteps=2, udd_callback=udd_callback) state = sim.get_state() assert np.array_equal(state.time, sim.data.time) assert np.array_equal(state.qpos, sim.data.qpos) assert np.array_equal(state.qvel, sim.data.qvel) assert np.array_equal(state.act, sim.data.act) for k in state.udd_state.keys(): if (isinstance(state.udd_state[k], Number)): assert state.udd_state[k] == sim.udd_state[k] else: assert np.array_equal(state.udd_state[k], sim.udd_state[k]) # test flatten, unflatten a = state.flatten() assert len(a) == (1 + sim.model.nq + sim.model.nv + sim.model.na + 8) state2 = MjSimState.from_flattened(a, sim) assert np.array_equal(state.time, sim.data.time) assert np.array_equal(state.qpos, sim.data.qpos) assert np.array_equal(state.qvel, sim.data.qvel) assert np.array_equal(state.act, sim.data.act) for k in state2.udd_state.keys(): if (isinstance(state2.udd_state[k], Number)): assert state2.udd_state[k] == sim.udd_state[k] else: assert np.array_equal(state2.udd_state[k], sim.udd_state[k]) assert state2 == state assert not state2 != state # test equality with deleting keys state2 = state2._replace(udd_state={"foo": foo}) assert state2 != state assert not (state2 == state) # test equality with changing contents of array state2 = state2._replace( udd_state={"foo": foo, "foo_array": np.array([foo, foo + 1])}) assert state2 != state assert not (state2 == state) # test equality with adding keys d2 = dict(d) d2.update({"not_foo": foo}) state2 = state2._replace(udd_state=d2) assert state2 != state assert not (state2 == state) # test defensive copy sim.set_state(state) state.qpos[0] = -1 assert not np.array_equal(state.qpos, sim.data.qpos) state3 = sim.get_state() state3.qpos[0] = -1 assert not np.array_equal(state3.qpos, sim.data.qpos) state3.udd_state["foo_array"][0] = -1 assert not np.array_equal( state3.udd_state["foo_array"], sim.udd_state["foo_array"]) # test no callback sim = MjSim(model, nsubsteps=2) state = sim.get_state() print("state.udd_state = %s" % state.udd_state) assert state.udd_state == {} # test flatten, unflatten a = state.flatten() assert len(a) == 1 + sim.model.nq + sim.model.nv + sim.model.na state2 = MjSimState.from_flattened(a, sim) assert np.array_equal(state.time, sim.data.time) assert np.array_equal(state.qpos, sim.data.qpos) assert np.array_equal(state.qvel, sim.data.qvel) assert np.array_equal(state.act, sim.data.act) assert state.udd_state == sim.udd_state
def test_inverse_kinematics(): panda_kinematics = ikpy_panda_kinematics.panda_kinematics() model = load_model_from_path("robot.xml") sim = MjSim(model) viewer = MjViewer(sim) timestep = generatePatternedTrajectories.TIMESTEP control_freq = 1/timestep total_time = 3 num_cycles = int(total_time * control_freq) qd_init = np.zeros(7) plt.ion() LW = 1.0 fig = plt.figure(figsize=(4,15)) axes = [] lines = [] goals = [] min_vals = {'x': -0.5, 'y': -0.5, 'z': 0.2} max_vals = {'x': 0.5, 'y': 0.5, 'z': 0.7} ylabels = ["x(m)", "y(m)", "z(m)"] ylbounds = [min_vals['x'], min_vals['y'], min_vals['z']] yubounds = [max_vals['x'], max_vals['y'], max_vals['z']] for i in range(3): axes.append(fig.add_subplot(3,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([ylbounds[i], yubounds[i]]) axes[i].set_xlim([0,total_time]) axes[i].set_ylabel(ylabels[i], fontsize=8) axes[i].set_xlabel("Time(s)", fontsize=8) for test in range(5): x_target = np.random.rand() * (max_vals['x'] - min_vals['x']) + min_vals['x'] y_target = np.random.rand() * (max_vals['y'] - min_vals['y']) + min_vals['y'] z_target = np.random.rand() * (max_vals['z'] - min_vals['z']) + min_vals['z'] if\ min_vals['x'] > x_target or max_vals['x'] < x_target or \ min_vals['y'] > y_target or max_vals['y'] < y_target or \ min_vals['z'] > z_target or max_vals['z'] < z_target: print("Error! 'xpos' out of range!") print("x = %.3f\ny = %.3f\nz = %.3f" % (x_target, y_target, z_target)) print(max_vals['y'] - min_vals['y']) return # x_target, y_target, z_target = 0.088, -0.0, 0.926 # roll = np.random.rand() * np.pi - np.pi/2 # pitch = np.random.rand() * np.pi - np.pi/2 # yaw = np.random.rand() * np.pi - np.pi/2 roll = 0 pitch = 0 yaw = np.pi ee_goal = [x_target, y_target, z_target, roll, pitch, yaw] # temp = bounds.getRandPosInBounds() # ee_goal = panda_kinematics.forward_kinematics(temp) q_init = bounds.getRandPosInBounds() q_goal = panda_kinematics.inverse_kinematics(translation=ee_goal[0:3], rpy=ee_goal[3:6], init_qpos=q_init) for g in range(3): goals[g].set_ydata(np.ones(num_cycles) * ee_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.PDControl(q=q,qd=qd,qgoal=q_goal) sim.step() viewer.render() if i % 70 == 0: xpos = panda_kinematics.forward_kinematics(q) for a in range(3): lines[a].set_xdata(np.append(lines[a].get_xdata(), sim_time)) lines[a].set_ydata(np.append(lines[a].get_ydata(), xpos[a])) fig.canvas.draw() fig.canvas.flush_events() print("[q\t]:{}".format(np.around(q,3))) print("[qgoal\t]:{}".format(np.around(q_goal,3))) print("[qgoal2\t]:{}".format(np.around(panda_kinematics.inverse_kinematics(ee_goal[0:3], ee_goal[3:6]),3))) print("[EE\t]:{}".format(np.around(panda_kinematics.forward_kinematics(q),3))) print("[EEgoal\t]:{}".format(np.around(ee_goal,3))) sim_time += timestep # panda_kinematics.plot_stick_figure(q) for i in range(3): lines[i].set_xdata([]) lines[i].set_ydata([]) time.sleep(1)
def set_state(s, env): env.reset() mj_state = MjSimState.from_flattened(s, env.sim) env.sim.set_state(mj_state) env.sim.forward() return env.env._get_obs()