Source

ars / demos / VehicleWithArm.py

Full commit
"""Runs a simulation of a vehicle with two powered wheels and one
free-rotating spherical wheel. It has a 2-link robotic arm attached,
with joints either friction-free or with friction proportional to
joint speed.

"""
from ars.app import Program, dispatcher, logger
import ars.utils.mathematical as mut
import ars.constants as cts
from ars.model.simulator import signals


class VehicleWithArm(Program):

	STEPS_PER_FRAME = 200
	BACKGROUND_COLOR = (0.8, 0.8, 0.8)
	FLOOR_BOX_SIZE = (20, 0.01, 20)

	WHEEL_TORQUE = 3
	VEHICLE_OFFSET = (2, 0.35, 4)

	# ((radius, center), density)
	BALL_PARAMS = ((0.3, (1, 0, 0)), {'density': 1})
	# ((size, center), mass)
	CHASSIS_PARAMS = (((2, 0.2, 1.5), (0.5, 0.45, 0)), {'mass': 6})

	# ((length, radius, center), density)
	WHEEL_R_PARAMS = ((0.4, 0.3, (0, 0, -0.5)), {'density': 1})
	WHEEL_L_PARAMS = ((0.4, 0.3, (0, 0, 0.5)), {'density': 1})

	# ((length, radius, center), density)
	LINK1_PARAMS = ((0.8, 0.1, (0, 0, 0)), {'density': 1})
	LINK2_PARAMS = ((0.6, 0.1, (0, 0.7, 0.2)), {'density': 1})

	R1_TORQUE = 3

	Q1_FRICTION_COEFF = 0.01
	Q2_FRICTION_COEFF = 0.01

	def __init__(self, use_capsule_wheels=False, frictionless_arm=True):
		"""Constructor, calls the superclass constructor first."""
		self._use_capsule_wheels = use_capsule_wheels
		self._frictionless_arm = frictionless_arm

		Program.__init__(self)
		self.key_press_functions.add('up', self.go_forwards, repeat=True)
		self.key_press_functions.add('down', self.go_backwards, repeat=True)
		self.key_press_functions.add('left', self.turn_left, repeat=True)
		self.key_press_functions.add('right', self.turn_right, repeat=True)

		self.key_press_functions.add('a', self.rotate_clockwise)
		self.key_press_functions.add('z', self.rotate_counterlockwise)

		dispatcher.connect(self.on_pre_step, signals.SIM_PRE_STEP)

	def create_sim_objects(self):
		"""Implementation of the required method.

		Creates and sets up all the objects of the simulation.

		"""
		arm_offset = (0, 0.5, 0)

		#=======================================================================
		# VEHICLE
		#=======================================================================

		if self._use_capsule_wheels:
			wheelR = self.sim.add_capsule(
				*self.WHEEL_R_PARAMS[0], **self.WHEEL_R_PARAMS[1])
			wheelL = self.sim.add_capsule(
				*self.WHEEL_L_PARAMS[0], **self.WHEEL_L_PARAMS[1])
		else:
			wheelR = self.sim.add_cylinder(
				*self.WHEEL_R_PARAMS[0], **self.WHEEL_R_PARAMS[1])
			wheelL = self.sim.add_cylinder(
				*self.WHEEL_L_PARAMS[0], **self.WHEEL_L_PARAMS[1])

		ball = self.sim.add_sphere(
			*self.BALL_PARAMS[0], **self.BALL_PARAMS[1])
		chassis = self.sim.add_box(
			*self.CHASSIS_PARAMS[0], **self.CHASSIS_PARAMS[1])

		# create joints: 2 rotary, 1 ball & socket
		self.sim.add_rotary_joint(
			'w1',                          # name
			self.sim.get_object(chassis),  # obj1
			self.sim.get_object(wheelR),   # obj2
			None,                          # anchor
			cts.Z_AXIS)                    # axis

		self.sim.add_rotary_joint(
			'w2',
			self.sim.get_object(chassis),
			self.sim.get_object(wheelL),
			None,
			cts.Z_AXIS)

		self.sim.add_ball_socket_joint(
			'bs',                          # name
			self.sim.get_object(chassis),  # obj1
			self.sim.get_object(ball),     # obj2
			None)                          # anchor

		self.sim.get_object(wheelR).offset_by_position(self.VEHICLE_OFFSET)
		self.sim.get_object(wheelL).offset_by_position(self.VEHICLE_OFFSET)
		self.sim.get_object(ball).offset_by_position(self.VEHICLE_OFFSET)
		self.sim.get_object(chassis).offset_by_position(self.VEHICLE_OFFSET)

		#=======================================================================
		# ROBOTIC ARM
		#=======================================================================
		link1 = self.sim.add_capsule(
			*self.LINK1_PARAMS[0], **self.LINK1_PARAMS[1])
		link2 = self.sim.add_capsule(
			*self.LINK2_PARAMS[0], **self.LINK2_PARAMS[1])

		# bodies are rotated before attaching themselves through joints
		self.sim.get_object(link1).rotate(cts.X_AXIS, mut.pi / 2)
		self.sim.get_object(link2).rotate(cts.X_AXIS, mut.pi / 2)

		self.sim.get_object(link1).offset_by_object(self.sim.get_object(chassis))
		self.sim.get_object(link1).offset_by_position(arm_offset)
		self.sim.get_object(link2).offset_by_object(self.sim.get_object(link1))

		self.sim.add_rotary_joint(
			'r1',
			self.sim.get_object(chassis),
			self.sim.get_object(link1),
			None,
			cts.Y_AXIS)

		r2_anchor = mut.sub3(
			self.sim.get_object(link2).get_position(),
			(0, self.LINK2_PARAMS[0][0] / 2, 0))  # (0, length/2, 0)
		self.sim.add_rotary_joint(
			'r2',
			self.sim.get_object(link1),
			self.sim.get_object(link2),
			r2_anchor,
			cts.Z_AXIS)

		try:
			self.sim.get_object(chassis).actor.set_color(cts.COLOR_RED)
			self.sim.get_object(link1).actor.set_color(cts.COLOR_YELLOW)
			self.sim.get_object(link2).actor.set_color(cts.COLOR_NAVY)
		except AttributeError:
			# if visualization is deactivated, there is no actor
			pass

		self.chassis = chassis
		self.wheelR = wheelR
		self.wheelL = wheelL
		self.ball = ball

		self.link1 = link1
		self.link2 = link2

	def go_forwards(self):
		"""Rotate both powered wheels in the same direction, forwards."""
		self.apply_torque_to_wheels(self.WHEEL_TORQUE, self.WHEEL_TORQUE)

	def go_backwards(self):
		"""Rotate both powered wheels in the same direction, backwards."""
		self.apply_torque_to_wheels(-self.WHEEL_TORQUE, -self.WHEEL_TORQUE)

	def turn_left(self):
		"""Rotate vehicle counter-clockwise (from above)."""
		self.apply_torque_to_wheels(-self.WHEEL_TORQUE, self.WHEEL_TORQUE)

	def turn_right(self):
		"""Rotate vehicle clockwise (from above)."""
		self.apply_torque_to_wheels(self.WHEEL_TORQUE, -self.WHEEL_TORQUE)

	def on_pre_step(self):
		#print(self.sim.get_object(self.chassis).get_position())
		try:
			#time = self.sim.sim_time
			q1p = self.get_q1p()
			q2p = self.get_q2p()

			if not self._frictionless_arm:
				self.apply_friction(q1p, q2p)

			print('q1p: %f, q2p: %f' % (q1p, q2p))

		except Exception:
			logger.exception("Exception when executing on_pre_step")

	def apply_torque_to_wheels(self, torque1, torque2):
		if torque1 is not None:
			self.sim.get_joint('w1').add_torque(torque1)
		if torque2 is not None:
			self.sim.get_joint('w2').add_torque(torque2)

	def rotate_clockwise(self):
		self.apply_torque_to_joints(self.R1_TORQUE, 0)

	def rotate_counterlockwise(self):
		self.apply_torque_to_joints(-self.R1_TORQUE, 0)

	def apply_torque_to_joints(self, torque1, torque2):
		if torque1 is not None:
			self.sim.get_joint('r1').add_torque(torque1)
		if torque2 is not None:
			self.sim.get_joint('r2').add_torque(torque2)

	def get_q1p(self):
		return self.sim.get_joint('r1').joint.angle_rate

	def get_q2p(self):
		return self.sim.get_joint('r2').joint.angle_rate

	def apply_friction(self, q1p, q2p):
		self.apply_torque_to_joints(
			-q1p * self.Q1_FRICTION_COEFF,
			-q2p * self.Q2_FRICTION_COEFF)