111 lines
4.4 KiB
Python
111 lines
4.4 KiB
Python
from .scene_stadium import SinglePlayerStadiumScene
|
|
from .env_bases import MJCFBaseBulletEnv
|
|
import numpy as np
|
|
from robot_locomotors import Hopper, Walker2D, HalfCheetah, Ant, Humanoid
|
|
|
|
|
|
class WalkerBaseBulletEnv(MJCFBaseBulletEnv):
|
|
def __init__(self, robot):
|
|
MJCFBaseBulletEnv.__init__(self, robot)
|
|
self.camera_x = 0
|
|
self.walk_target_x = 1e3 # kilometer away
|
|
self.walk_target_y = 0
|
|
|
|
def create_single_player_scene(self):
|
|
self.stadium_scene = SinglePlayerStadiumScene(gravity=9.8, timestep=0.0165/4, frame_skip=4)
|
|
return self.stadium_scene
|
|
|
|
def _reset(self):
|
|
r = MJCFBaseBulletEnv._reset(self)
|
|
self.parts, self.jdict, self.ordered_joints, self.robot_body = self.robot.addToScene(
|
|
self.stadium_scene.ground_plane_mjcf)
|
|
self.ground_ids = set([(self.parts[f].bodies[self.parts[f].bodyIndex], self.parts[f].bodyPartIndex) for f in
|
|
self.foot_ground_object_names])
|
|
return r
|
|
|
|
def move_robot(self, init_x, init_y, init_z):
|
|
"Used by multiplayer stadium to move sideways, to another running lane."
|
|
self.cpp_robot.query_position()
|
|
pose = self.cpp_robot.root_part.pose()
|
|
pose.move_xyz(init_x, init_y, init_z) # Works because robot loads around (0,0,0), and some robots have z != 0 that is left intact
|
|
self.cpp_robot.set_pose(pose)
|
|
|
|
electricity_cost = -2.0 # cost for using motors -- this parameter should be carefully tuned against reward for making progress, other values less improtant
|
|
stall_torque_cost = -0.1 # cost for running electric current through a motor even at zero rotational speed, small
|
|
foot_collision_cost = -1.0 # touches another leg, or other objects, that cost makes robot avoid smashing feet into itself
|
|
foot_ground_object_names = set(["floor"]) # to distinguish ground and other objects
|
|
joints_at_limit_cost = -0.1 # discourage stuck joints
|
|
|
|
def _step(self, a):
|
|
if not self.scene.multiplayer: # if multiplayer, action first applied to all robots, then global step() called, then _step() for all robots with the same actions
|
|
self.robot.apply_action(a)
|
|
self.scene.global_step()
|
|
|
|
state = self.robot.calc_state() # also calculates self.joints_at_limit
|
|
|
|
alive = float(self.robot.alive_bonus(state[0]+self.robot.initial_z, self.robot.body_rpy[1])) # state[0] is body height above ground, body_rpy[1] is pitch
|
|
done = alive < 0
|
|
if not np.isfinite(state).all():
|
|
print("~INF~", state)
|
|
done = True
|
|
|
|
potential_old = self.potential
|
|
self.potential = self.robot.calc_potential()
|
|
progress = float(self.potential - potential_old)
|
|
|
|
feet_collision_cost = 0.0
|
|
for i,f in enumerate(self.robot.feet): # TODO: Maybe calculating feet contacts could be done within the robot code
|
|
contact_ids = set((x[2], x[4]) for x in f.contact_list())
|
|
#print("CONTACT OF '%s' WITH %s" % (f.name, ",".join(contact_names)) )
|
|
self.robot.feet_contact[i] = 1.0 if (self.ground_ids & contact_ids) else 0.0
|
|
if contact_ids - self.ground_ids:
|
|
feet_collision_cost += self.foot_collision_cost
|
|
|
|
electricity_cost = self.electricity_cost * float(np.abs(a*self.robot.joint_speeds).mean()) # let's assume we have DC motor with controller, and reverse current braking
|
|
electricity_cost += self.stall_torque_cost * float(np.square(a).mean())
|
|
|
|
joints_at_limit_cost = float(self.joints_at_limit_cost * self.robot.joints_at_limit)
|
|
|
|
self.rewards = [
|
|
alive,
|
|
progress,
|
|
electricity_cost,
|
|
joints_at_limit_cost,
|
|
feet_collision_cost
|
|
]
|
|
|
|
self.HUD(state, a, done)
|
|
return state, sum(self.rewards), bool(done), {}
|
|
|
|
def camera_adjust(self):
|
|
x, y, z = self.body_xyz
|
|
self.camera_x = 0.98*self.camera_x + (1-0.98)*x
|
|
self.camera.move_and_look_at(self.camera_x, y-2.0, 1.4, x, y, 1.0)
|
|
|
|
class HopperBulletEnv(WalkerBaseBulletEnv):
|
|
def __init__(self):
|
|
self.robot = Hopper()
|
|
WalkerBaseBulletEnv.__init__(self, self.robot)
|
|
|
|
class Walker2DBulletEnv(WalkerBaseBulletEnv):
|
|
def __init__(self):
|
|
self.robot = Walker2D()
|
|
WalkerBaseBulletEnv.__init__(self, self.robot)
|
|
|
|
class HalfCheetahBulletEnv(WalkerBaseBulletEnv):
|
|
def __init__(self):
|
|
self.robot = HalfCheetah()
|
|
WalkerBaseBulletEnv.__init__(self, self.robot)
|
|
|
|
class AntBulletEnv(WalkerBaseBulletEnv):
|
|
def __init__(self):
|
|
self.robot = Ant()
|
|
WalkerBaseBulletEnv.__init__(self, self.robot)
|
|
|
|
class HumanoidBulletEnv(WalkerBaseBulletEnv):
|
|
def __init__(self):
|
|
self.robot = Humanoid()
|
|
WalkerBaseBulletEnv.__init__(self, self.robot)
|
|
self.electricity_cost = 4.25*WalkerBaseBulletEnv.electricity_cost
|
|
self.stall_torque_cost = 4.25*WalkerBaseBulletEnv.stall_torque_cost
|