from robot_bases import MJCFBasedRobot
import numpy as np
class InvertedPendulum(MJCFBasedRobot):
swingup = False
def __init__(self):
MJCFBasedRobot.__init__(self, 'inverted_pendulum.xml', 'cart', action_dim=1, obs_dim=5)
def robot_specific_reset(self, bullet_client):
self._p = bullet_client
self.pole = self.parts["pole"]
self.slider = self.jdict["slider"]
self.j1 = self.jdict["hinge"]
u = self.np_random.uniform(low=-.1, high=.1)
self.j1.reset_current_position( u if not self.swingup else 3.1415+u , 0)
self.j1.set_motor_torque(0)
def apply_action(self, a):
assert( np.isfinite(a).all() )
if not np.isfinite(a).all():
print("a is inf")
a[0] = 0
self.slider.set_motor_torque( 100*float(np.clip(a[0], -1, +1)) )
def calc_state(self):
self.theta, theta_dot = self.j1.current_position()
x, vx = self.slider.current_position()
assert( np.isfinite(x) )
if not np.isfinite(x):
print("x is inf")
x = 0
if not np.isfinite(vx):
print("vx is inf")
vx = 0
if not np.isfinite(self.theta):
print("theta is inf")
self.theta = 0
if not np.isfinite(theta_dot):
print("theta_dot is inf")
theta_dot = 0
return np.array([
x, vx,
np.cos(self.theta), np.sin(self.theta), theta_dot
])
class InvertedPendulumSwingup(InvertedPendulum):
swingup = True
class InvertedDoublePendulum(MJCFBasedRobot):
def __init__(self):
MJCFBasedRobot.__init__(self, 'inverted_double_pendulum.xml', 'cart', action_dim=1, obs_dim=9)
def robot_specific_reset(self, bullet_client):
self._p = bullet_client
self.pole2 = self.parts["pole2"]
self.slider = self.jdict["slider"]
self.j1 = self.jdict["hinge"]
self.j2 = self.jdict["hinge2"]
u = self.np_random.uniform(low=-.1, high=.1, size=[2])
self.j1.reset_current_position(float(u[0]), 0)
self.j2.reset_current_position(float(u[1]), 0)
self.j1.set_motor_torque(0)
self.j2.set_motor_torque(0)
def apply_action(self, a):
assert( np.isfinite(a).all() )
self.slider.set_motor_torque( 200*float(np.clip(a[0], -1, +1)) )
def calc_state(self):
theta, theta_dot = self.j1.current_position()
gamma, gamma_dot = self.j2.current_position()
x, vx = self.slider.current_position()
self.pos_x, _, self.pos_y = self.pole2.pose().xyz()
assert( np.isfinite(x) )
return np.array([
x, vx,
self.pos_x,
np.cos(theta), np.sin(theta), theta_dot,
np.cos(gamma), np.sin(gamma), gamma_dot,
])
