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
