-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathprecession.py
58 lines (46 loc) · 1.2 KB
/
precession.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
import xpdb_brax
from xpdb_brax.physics.config import Config, Box, Plane, Joint
from xpdb_brax.viewer.server import MyViewer
from xpdb_brax.physics.base import QP
from xpdb_brax.physics.system import System
import numpy as np
import jax.numpy as jnp
import jax
import time
config = Config()
qp = []
config.bodies.append(Box(1, [.1, .1, .1], frozen=True))
qp.append(QP(
jnp.array([0., 0, 1]),
jnp.array([1., 0, 0, 0]),
jnp.array([0., 0, 0]),
jnp.array([0., 0, 0]),
))
config.bodies.append(Box(1, [.05, 2, 2]))
qp.append(QP(
jnp.array([-0.5, 0, 1]),
jnp.array([1., 0, 0, 0]),
jnp.array([0, 0, 0]),
jnp.array([-10., 0, 0]),
))
config.joints.append(Joint ("ball", 0, 1, [0., 0, 0], [0.5, 0, 0], [1, 0, 0], 0))
config.gravity = [0, 0, -9.81]
config.dt = 0.01
config.substeps = 1
qp = jax.tree_multimap((lambda *args: jnp.stack(args)), *qp)
# print(type(a))
# exit()
system = System(config)
render = True
if render:
xpdb_viewer = MyViewer()
xpdb_viewer.init(config)
for i in range(3 if not render else 1000):
qp = system.step(qp)
pos, rot = qp.pos, qp.rot
if render:
xpdb_viewer.update(pos, rot)
time.sleep(0.001)
if render:
input()
print("success !!")