mirror of
https://github.com/illiumst/marl-factory-grid.git
synced 2025-10-31 04:37:25 +01:00
290 lines
11 KiB
Python
290 lines
11 KiB
Python
"""This file is based on the car_dynamics.py in OpenAI's gym.
|
|
|
|
Description
|
|
We provide some modified CarRacing environments here:
|
|
* CarRacingColor-v0
|
|
Both the lane and the grass colors are perturbed at the env.reset() with scalar noises.
|
|
* CarRacingColor3-v0
|
|
Both the lane and the grass colors are perturbed at the env.reset() with 3d noises.
|
|
* CarRacingBar-v0
|
|
We add vertical bars on the left and right side of the screen.
|
|
* CarRacingBlob-v0
|
|
A red blob that follows the car at a fixed position in the car's frame.
|
|
* CarRacingNoise-v0
|
|
We replace the green background with noise.
|
|
* CarRacingVideo-v0
|
|
We replace the green background with frames from a video, the user is
|
|
responsible for creating these frames and pass the directory.
|
|
These environments were first used in the paper "Neuroevolution of Self-Interpretable Agent"
|
|
https://attentionagent.github.io/
|
|
|
|
Author
|
|
Yujin Tang (yujintang@google.com)
|
|
"""
|
|
|
|
import numpy as np
|
|
import math
|
|
import Box2D
|
|
from Box2D.b2 import (edgeShape, circleShape, fixtureDef, polygonShape, revoluteJointDef, contactListener, shape)
|
|
|
|
# Top-down car dynamics simulation.
|
|
#
|
|
# Some ideas are taken from this great tutorial http://www.iforce2d.net/b2dtut/top-down-car by Chris Campbell.
|
|
# This simulation is a bit more detailed, with wheels rotation.
|
|
#
|
|
# Created by Oleg Klimov. Licensed on the same terms as the rest of OpenAI Gym.
|
|
|
|
SIZE = 0.02
|
|
ENGINE_POWER = 100000000*SIZE*SIZE
|
|
WHEEL_MOMENT_OF_INERTIA = 4000*SIZE*SIZE
|
|
FRICTION_LIMIT = 1000000*SIZE*SIZE # friction ~= mass ~= size^2 (calculated implicitly using density)
|
|
WHEEL_R = 27
|
|
WHEEL_W = 14
|
|
WHEELPOS = [
|
|
(-55,+80), (+55,+80),
|
|
(-55,-82), (+55,-82)
|
|
]
|
|
HULL_POLY1 =[
|
|
(-60,+130), (+60,+130),
|
|
(+60,+110), (-60,+110)
|
|
]
|
|
HULL_POLY2 =[
|
|
(-15,+120), (+15,+120),
|
|
(+20, +20), (-20, 20)
|
|
]
|
|
HULL_POLY3 =[
|
|
(+25, +20),
|
|
(+50, -10),
|
|
(+50, -40),
|
|
(+20, -90),
|
|
(-20, -90),
|
|
(-50, -40),
|
|
(-50, -10),
|
|
(-25, +20)
|
|
]
|
|
HULL_POLY4 =[
|
|
(-50,-120), (+50,-120),
|
|
(+50,-90), (-50,-90)
|
|
]
|
|
|
|
WHEEL_COLOR = (0.0,0.0,0.0)
|
|
WHEEL_WHITE = (0.3,0.3,0.3)
|
|
MUD_COLOR = (0.4,0.4,0.0)
|
|
|
|
class Car:
|
|
def __init__(self, world, init_angle, init_x, init_y, add_blob=False):
|
|
self.world = world
|
|
|
|
self.add_blob = add_blob
|
|
|
|
fixtures = [
|
|
fixtureDef(shape=polygonShape(
|
|
vertices=[ (x*SIZE,y*SIZE) for x,y in HULL_POLY1 ]), density=1.0),
|
|
fixtureDef(shape=polygonShape(
|
|
vertices=[ (x*SIZE,y*SIZE) for x,y in HULL_POLY2 ]), density=1.0),
|
|
fixtureDef(shape=polygonShape(
|
|
vertices=[ (x*SIZE,y*SIZE) for x,y in HULL_POLY3 ]), density=1.0),
|
|
fixtureDef(shape=polygonShape(
|
|
vertices=[ (x*SIZE,y*SIZE) for x,y in HULL_POLY4 ]), density=1.0),
|
|
]
|
|
self.hull = self.world.CreateDynamicBody(
|
|
position = (init_x, init_y),
|
|
angle = init_angle,
|
|
fixtures = fixtures,
|
|
)
|
|
self.hull.color = (0.8,0.0,0.0)
|
|
self.wheels = []
|
|
self.fuel_spent = 0.0
|
|
WHEEL_POLY = [
|
|
(-WHEEL_W,+WHEEL_R), (+WHEEL_W,+WHEEL_R),
|
|
(+WHEEL_W,-WHEEL_R), (-WHEEL_W,-WHEEL_R)
|
|
]
|
|
for wx,wy in WHEELPOS:
|
|
front_k = 1.0 if wy > 0 else 1.0
|
|
w = self.world.CreateDynamicBody(
|
|
position = (init_x+wx*SIZE, init_y+wy*SIZE),
|
|
angle = init_angle,
|
|
fixtures = fixtureDef(
|
|
shape=polygonShape(vertices=[ (x*front_k*SIZE,y*front_k*SIZE) for x,y in WHEEL_POLY ]),
|
|
density=0.1,
|
|
categoryBits=0x0020,
|
|
maskBits=0x001,
|
|
restitution=0.0)
|
|
)
|
|
w.wheel_rad = front_k*WHEEL_R*SIZE
|
|
w.color = WHEEL_COLOR
|
|
w.gas = 0.0
|
|
w.brake = 0.0
|
|
w.steer = 0.0
|
|
w.phase = 0.0 # wheel angle
|
|
w.omega = 0.0 # angular velocity
|
|
w.skid_start = None
|
|
w.skid_particle = None
|
|
rjd = revoluteJointDef(
|
|
bodyA=self.hull,
|
|
bodyB=w,
|
|
localAnchorA=(wx*SIZE,wy*SIZE),
|
|
localAnchorB=(0,0),
|
|
enableMotor=True,
|
|
enableLimit=True,
|
|
maxMotorTorque=180*900*SIZE*SIZE,
|
|
motorSpeed = 0,
|
|
lowerAngle = -0.4,
|
|
upperAngle = +0.4,
|
|
)
|
|
w.joint = self.world.CreateJoint(rjd)
|
|
w.tiles = set()
|
|
w.userData = w
|
|
self.wheels.append(w)
|
|
self.drawlist = self.wheels + [self.hull]
|
|
self.particles = []
|
|
|
|
def gas(self, gas):
|
|
'control: rear wheel drive'
|
|
gas = np.clip(gas, 0, 1)
|
|
for w in self.wheels[2:4]:
|
|
diff = gas - w.gas
|
|
if diff > 0.1: diff = 0.1 # gradually increase, but stop immediately
|
|
w.gas += diff
|
|
|
|
def brake(self, b):
|
|
'control: brake b=0..1, more than 0.9 blocks wheels to zero rotation'
|
|
for w in self.wheels:
|
|
w.brake = b
|
|
|
|
def steer(self, s):
|
|
'control: steer s=-1..1, it takes time to rotate steering wheel from side to side, s is target position'
|
|
self.wheels[0].steer = s
|
|
self.wheels[1].steer = s
|
|
|
|
def step(self, dt):
|
|
for w in self.wheels:
|
|
# Steer each wheel
|
|
dir = np.sign(w.steer - w.joint.angle)
|
|
val = abs(w.steer - w.joint.angle)
|
|
w.joint.motorSpeed = dir*min(50.0*val, 3.0)
|
|
|
|
# Position => friction_limit
|
|
grass = True
|
|
friction_limit = FRICTION_LIMIT*0.6 # Grass friction if no tile
|
|
for tile in w.tiles:
|
|
friction_limit = max(friction_limit, FRICTION_LIMIT*tile.road_friction)
|
|
grass = False
|
|
|
|
# Force
|
|
forw = w.GetWorldVector( (0,1) )
|
|
side = w.GetWorldVector( (1,0) )
|
|
v = w.linearVelocity
|
|
vf = forw[0]*v[0] + forw[1]*v[1] # forward speed
|
|
vs = side[0]*v[0] + side[1]*v[1] # side speed
|
|
|
|
# WHEEL_MOMENT_OF_INERTIA*np.square(w.omega)/2 = E -- energy
|
|
# WHEEL_MOMENT_OF_INERTIA*w.omega * domega/dt = dE/dt = W -- power
|
|
# domega = dt*W/WHEEL_MOMENT_OF_INERTIA/w.omega
|
|
w.omega += dt*ENGINE_POWER*w.gas/WHEEL_MOMENT_OF_INERTIA/(abs(w.omega)+5.0) # small coef not to divide by zero
|
|
self.fuel_spent += dt*ENGINE_POWER*w.gas
|
|
|
|
if w.brake >= 0.9:
|
|
w.omega = 0
|
|
elif w.brake > 0:
|
|
BRAKE_FORCE = 15 # radians per second
|
|
dir = -np.sign(w.omega)
|
|
val = BRAKE_FORCE*w.brake
|
|
if abs(val) > abs(w.omega): val = abs(w.omega) # low speed => same as = 0
|
|
w.omega += dir*val
|
|
w.phase += w.omega*dt
|
|
|
|
vr = w.omega*w.wheel_rad # rotating wheel speed
|
|
f_force = -vf + vr # force direction is direction of speed difference
|
|
p_force = -vs
|
|
|
|
# Physically correct is to always apply friction_limit until speed is equal.
|
|
# But dt is finite, that will lead to oscillations if difference is already near zero.
|
|
f_force *= 205000*SIZE*SIZE # Random coefficient to cut oscillations in few steps (have no effect on friction_limit)
|
|
p_force *= 205000*SIZE*SIZE
|
|
force = np.sqrt(np.square(f_force) + np.square(p_force))
|
|
|
|
# Skid trace
|
|
if abs(force) > 2.0*friction_limit:
|
|
if w.skid_particle and w.skid_particle.grass==grass and len(w.skid_particle.poly) < 30:
|
|
w.skid_particle.poly.append( (w.position[0], w.position[1]) )
|
|
elif w.skid_start is None:
|
|
w.skid_start = w.position
|
|
else:
|
|
w.skid_particle = self._create_particle( w.skid_start, w.position, grass )
|
|
w.skid_start = None
|
|
else:
|
|
w.skid_start = None
|
|
w.skid_particle = None
|
|
|
|
if abs(force) > friction_limit:
|
|
f_force /= force
|
|
p_force /= force
|
|
force = friction_limit # Correct physics here
|
|
f_force *= force
|
|
p_force *= force
|
|
|
|
w.omega -= dt*f_force*w.wheel_rad/WHEEL_MOMENT_OF_INERTIA
|
|
|
|
w.ApplyForceToCenter( (
|
|
p_force*side[0] + f_force*forw[0],
|
|
p_force*side[1] + f_force*forw[1]), True )
|
|
|
|
def draw(self, viewer, draw_particles=True):
|
|
if draw_particles:
|
|
for p in self.particles:
|
|
viewer.draw_polyline(p.poly, color=p.color, linewidth=5)
|
|
for obj in self.drawlist:
|
|
for i, f in enumerate(obj.fixtures):
|
|
trans = f.body.transform
|
|
path = [trans*v for v in f.shape.vertices]
|
|
if self.add_blob and i == 3:
|
|
obj_color = (0.8, 0.0, 0.)
|
|
offset_x = f.shape.vertices[0][0] + 20
|
|
offset_y = f.shape.vertices[0][1] + 10
|
|
width = height = 8
|
|
blob =[
|
|
trans * (+offset_x, +offset_y+height),
|
|
trans * (+offset_x+width, +offset_y+height),
|
|
trans * (+offset_x+width, +offset_y),
|
|
trans * (+offset_x, +offset_y),
|
|
]
|
|
viewer.draw_polygon(blob, color=obj_color)
|
|
viewer.draw_polygon(path, color=obj.color)
|
|
if "phase" not in obj.__dict__: continue
|
|
a1 = obj.phase
|
|
a2 = obj.phase + 1.2 # radians
|
|
s1 = math.sin(a1)
|
|
s2 = math.sin(a2)
|
|
c1 = math.cos(a1)
|
|
c2 = math.cos(a2)
|
|
if s1>0 and s2>0: continue
|
|
if s1>0: c1 = np.sign(c1)
|
|
if s2>0: c2 = np.sign(c2)
|
|
white_poly = [
|
|
(-WHEEL_W*SIZE, +WHEEL_R*c1*SIZE), (+WHEEL_W*SIZE, +WHEEL_R*c1*SIZE),
|
|
(+WHEEL_W*SIZE, +WHEEL_R*c2*SIZE), (-WHEEL_W*SIZE, +WHEEL_R*c2*SIZE)
|
|
]
|
|
viewer.draw_polygon([trans*v for v in white_poly], color=WHEEL_WHITE)
|
|
|
|
def _create_particle(self, point1, point2, grass):
|
|
class Particle:
|
|
pass
|
|
p = Particle()
|
|
p.color = WHEEL_COLOR if not grass else MUD_COLOR
|
|
p.ttl = 1
|
|
p.poly = [(point1[0],point1[1]), (point2[0],point2[1])]
|
|
p.grass = grass
|
|
self.particles.append(p)
|
|
while len(self.particles) > 30:
|
|
self.particles.pop(0)
|
|
return p
|
|
|
|
def destroy(self):
|
|
self.world.DestroyBody(self.hull)
|
|
self.hull = None
|
|
for w in self.wheels:
|
|
self.world.DestroyBody(w)
|
|
self.wheels = []
|
|
|