Files
bullet3/examples/pybullet/examples/snake.py
Erwin Coumans ef9570c315 add yapf style and apply yapf to format all Python files
This recreates pull request #2192
2019-04-27 07:31:15 -07:00

143 lines
4.7 KiB
Python

import pybullet as p
import time
import math
# This simple snake logic is from some 15 year old Havok C++ demo
# Thanks to Michael Ewert!
p.connect(p.GUI)
plane = p.createCollisionShape(p.GEOM_PLANE)
p.createMultiBody(0, plane)
useMaximalCoordinates = True
sphereRadius = 0.25
#colBoxId = p.createCollisionShapeArray([p.GEOM_BOX, p.GEOM_SPHERE],radii=[sphereRadius+0.03,sphereRadius+0.03], halfExtents=[[sphereRadius,sphereRadius,sphereRadius],[sphereRadius,sphereRadius,sphereRadius]])
colBoxId = p.createCollisionShape(p.GEOM_BOX,
halfExtents=[sphereRadius, sphereRadius, sphereRadius])
mass = 1
visualShapeId = -1
link_Masses = []
linkCollisionShapeIndices = []
linkVisualShapeIndices = []
linkPositions = []
linkOrientations = []
linkInertialFramePositions = []
linkInertialFrameOrientations = []
indices = []
jointTypes = []
axis = []
for i in range(36):
link_Masses.append(1)
linkCollisionShapeIndices.append(colBoxId)
linkVisualShapeIndices.append(-1)
linkPositions.append([0, sphereRadius * 2.0 + 0.01, 0])
linkOrientations.append([0, 0, 0, 1])
linkInertialFramePositions.append([0, 0, 0])
linkInertialFrameOrientations.append([0, 0, 0, 1])
indices.append(i)
jointTypes.append(p.JOINT_REVOLUTE)
axis.append([0, 0, 1])
basePosition = [0, 0, 1]
baseOrientation = [0, 0, 0, 1]
sphereUid = p.createMultiBody(mass,
colBoxId,
visualShapeId,
basePosition,
baseOrientation,
linkMasses=link_Masses,
linkCollisionShapeIndices=linkCollisionShapeIndices,
linkVisualShapeIndices=linkVisualShapeIndices,
linkPositions=linkPositions,
linkOrientations=linkOrientations,
linkInertialFramePositions=linkInertialFramePositions,
linkInertialFrameOrientations=linkInertialFrameOrientations,
linkParentIndices=indices,
linkJointTypes=jointTypes,
linkJointAxis=axis,
useMaximalCoordinates=useMaximalCoordinates)
p.setGravity(0, 0, -10)
p.setRealTimeSimulation(0)
anistropicFriction = [1, 0.01, 0.01]
p.changeDynamics(sphereUid, -1, lateralFriction=2, anisotropicFriction=anistropicFriction)
p.getNumJoints(sphereUid)
for i in range(p.getNumJoints(sphereUid)):
p.getJointInfo(sphereUid, i)
p.changeDynamics(sphereUid, i, lateralFriction=2, anisotropicFriction=anistropicFriction)
dt = 1. / 240.
SNAKE_NORMAL_PERIOD = 0.1 #1.5
m_wavePeriod = SNAKE_NORMAL_PERIOD
m_waveLength = 4
m_wavePeriod = 1.5
m_waveAmplitude = 0.4
m_waveFront = 0.0
#our steering value
m_steering = 0.0
m_segmentLength = sphereRadius * 2.0
forward = 0
while (1):
keys = p.getKeyboardEvents()
for k, v in keys.items():
if (k == p.B3G_RIGHT_ARROW and (v & p.KEY_WAS_TRIGGERED)):
m_steering = -.2
if (k == p.B3G_RIGHT_ARROW and (v & p.KEY_WAS_RELEASED)):
m_steering = 0
if (k == p.B3G_LEFT_ARROW and (v & p.KEY_WAS_TRIGGERED)):
m_steering = .2
if (k == p.B3G_LEFT_ARROW and (v & p.KEY_WAS_RELEASED)):
m_steering = 0
amp = 0.2
offset = 0.6
numMuscles = p.getNumJoints(sphereUid)
scaleStart = 1.0
#start of the snake with smaller waves.
#I think starting the wave at the tail would work better ( while it still goes from head to tail )
if (m_waveFront < m_segmentLength * 4.0):
scaleStart = m_waveFront / (m_segmentLength * 4.0)
segment = numMuscles - 1
#we simply move a sin wave down the body of the snake.
#this snake may be going backwards, but who can tell ;)
for joint in range(p.getNumJoints(sphereUid)):
segment = joint #numMuscles-1-joint
#map segment to phase
phase = (m_waveFront - (segment + 1) * m_segmentLength) / m_waveLength
phase -= math.floor(phase)
phase *= math.pi * 2.0
#map phase to curvature
targetPos = math.sin(phase) * scaleStart * m_waveAmplitude
#// steer snake by squashing +ve or -ve side of sin curve
if (m_steering > 0 and targetPos < 0):
targetPos *= 1.0 / (1.0 + m_steering)
if (m_steering < 0 and targetPos > 0):
targetPos *= 1.0 / (1.0 - m_steering)
#set our motor
p.setJointMotorControl2(sphereUid,
joint,
p.POSITION_CONTROL,
targetPosition=targetPos + m_steering,
force=30)
#wave keeps track of where the wave is in time
m_waveFront += dt / m_wavePeriod * m_waveLength
p.stepSimulation()
time.sleep(dt)