def pendulum_loss(mx,
Sx,
target=np.array([np.pi, 0]),
angle_dims=[0],
pole_length=0.5,
cw=[0.25],
*args,
**kwargs):
# size of target vector (and mx) after replacing angles with their
# (sin, cos) representation:
# [x1,x2,..., angle,...,xn] -> [x1,x2,...,xn, sin(angle), cos(angle)]
Da = np.array(target).size + len(angle_dims)
# build cost scaling function
Q = np.zeros((Da, Da))
Q[0, 0] = 1
Q[0, -2] = pole_length
Q[-2, 0] = pole_length
Q[-2, -2] = pole_length**2
Q[-1, -1] = pole_length**2
return cost.distance_based_cost(mx,
Sx,
target,
Q,
cw,
angle_dims=angle_dims,
*args,
**kwargs)
def double_cartpole_loss(mx,
Sx,
target=np.array([0, 0, 0, 0, 0, 0]),
angle_dims=[4, 5],
link1_length=0.5,
link2_length=0.5,
cw=[0.5],
*args,
**kwargs):
# size of target vector (and mx) after replacing angles with their
# (sin, cos) representation:
# [x1,x2,..., angle,...,xn] -> [x1,x2,...,xn, sin(angle), cos(angle)]
Da = np.array(target).size + len(angle_dims)
# build cost scaling function
Q = np.zeros((Da, Da))
# these are the dimensions used to compute the cost
# (x, sin(theta1), cos(theta1), sin(theta2), cos(theta2))
cost_dims = np.hstack([0, np.arange(Da - 2 * len(angle_dims), Da)])[:,
None]
C = np.array([[1, -link1_length, 0, -link2_length, 0],
[0, 0, link1_length, 0, link2_length]])
Q[cost_dims, cost_dims.T] = C.T.dot(C)
return cost.distance_based_cost(mx,
Sx,
target,
Q,
cw,
angle_dims=angle_dims,
*args,
**kwargs)