-
Notifications
You must be signed in to change notification settings - Fork 1
/
quad.py
163 lines (133 loc) · 5.31 KB
/
quad.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
from random import random
from ode import (
World, GeomPlane, Space,
JointGroup, collide, ContactJoint
)
from visual import rate, vector, norm, degrees, box
from objs import Sphere
class Quad(object):
def __init__(self, world, space, pos=(0, 0, 0), size=1, motors_force=(0, 0, 0, 0)):
self.world = world
self.pos = pos
x, y, z = pos
self.size = size
self.motors_force = motors_force
self.motors_noise = [random()/100 for m in range(4)]
# Create the quad inside the world
# Motors
motor_size = size / 10.
motor_weight = 0.07
arms_distance = size / 2.
motor_front_color = (0, 0, 1) # color blue
motor_roll_color = (0, 1, 0) # color green
self.motors = [
Sphere(world, space, (x, y, z-arms_distance), motor_size, motor_weight, motor_front_color),
Sphere(world, space, (x, y, z+arms_distance), motor_size, motor_weight),
Sphere(world, space, (x-arms_distance, y, z), motor_size, motor_weight, motor_roll_color),
Sphere(world, space, (x+arms_distance, y, z), motor_size, motor_weight),
]
# Center
center_size = motor_size * 1.5
center_weight = 0.5
self.center = Sphere(world, space, (x, y, z), center_size, center_weight)
# Stick parts together
self.motors[0].joint(self.motors[1])
self.motors[0].joint(self.center)
self.motors[2].joint(self.motors[3])
self.motors[2].joint(self.center)
self.motors[0].joint(self.motors[2])
self.motors[1].joint(self.motors[3])
# set util vectors
self._set_plane_vectors()
def draw(self):
# draw motors
[m.draw() for m in self.motors]
self.center.draw()
def step(self):
# update force
self.update_motors()
# update pos
self.pos = self.center.getPosition()
# draw quad
self.draw()
def _set_plane_vectors(self):
cx, cy, cz = self.center.getPosition()
x, y, z = self.motors[0].getPosition()
self.v1 = norm(vector(x-cx, y-cy, z-cz))
x, y, z = self.motors[2].getPosition()
self.v2 = norm(vector(x-cx, y-cy, z-cz))
# Calculate pitch and roll
self.pitch = 90 - degrees(self.v1.diff_angle(vector(0, 1, 0)))
self.roll = degrees(self.v2.diff_angle(vector(0, 1, 0))) - 90
# calculate yaw, (this works if the quad is close to a horizontal position)
if self.v1.z <= 0:
self.yaw = degrees(self.v1.diff_angle(vector(1, 0, 0)))
else:
self.yaw = 360 - degrees(self.v1.diff_angle(vector(1, 0, 0)))
def update_motors(self):
# apply noise to motors to emulate reality
m_noise = self.motors_noise
# force of motor should be orthogonal to the quad
self._set_plane_vectors()
v = self.v1.cross(self.v2)
for i, f in enumerate(self.motors_force):
self.motors[i].addForce(v*(f+m_noise[i]))
# Add rotational fore to motors
rotational_factor = 0.1
rforce = self.v2*(self.motors_force[0]+m_noise[0])*rotational_factor
self.motors[0].addForce(rforce)
rforce = -self.v2*(self.motors_force[1]+m_noise[1])*rotational_factor
self.motors[1].addForce(rforce)
rforce = self.v1*(self.motors_force[2]+m_noise[2])*rotational_factor
self.motors[2].addForce(rforce)
rforce = -self.v1*(self.motors_force[3]+m_noise[3])*rotational_factor
self.motors[3].addForce(rforce)
@property
def vel(self):
return self.center.getLinearVel()
######################################
# Simulation test
if __name__ == '__main__':
world = World()
world.setGravity((0, -9.81, 0))
world.setERP(0.8)
world.setCFM(1E-5)
space = Space()
contactgroup = JointGroup()
# Set environment
floor = GeomPlane(space, (0, 1, 0), -1)
box(pos=(0, -1, 0), width=2, height=0.01, length=2)
quad = Quad(world, space, pos=(0, -0.5, 0))
# set initial state of motors
# quad.motors_noise = (0, 0, 0, 0)
quad.motors_force = (1.92, 1.92, 1.92, 1.92)
# Collision callback
def near_callback(args, geom1, geom2):
"""Callback function for the collide() method.
This function checks if the given geoms do collide and
creates contact joints if they do.
"""
# Check if the objects do collide
contacts = collide(geom1, geom2)
# Create contact joints
world, contactgroup = args
for c in contacts:
c.setBounce(0.2)
c.setMu(5000)
j = ContactJoint(world, contactgroup, c)
j.attach(geom1.getBody(), geom2.getBody())
# Do the simulation...
total_time = 0.0
dt = 0.04
while True:
rate(40)
print "%1.2fsec: pos=(%6.3f, %6.3f, %6.3f) pry=(%6.3f, %6.3f, %6.3f)" % \
(total_time, quad.pos[0], quad.pos[1], quad.pos[2],
quad.pitch, quad.roll, quad.yaw)
print "power=(%6.3f, %6.3f, %6.3f, %6.3f)" % (quad.motors_force[0], quad.motors_force[1], quad.motors_force[2], quad.motors_force[3])
# Detect collisions and create contact joints
space.collide((world, contactgroup), near_callback)
quad.step()
world.step(dt)
contactgroup.empty()
total_time+=dt