Weird pendulum behavior in Brax Basics notebook

[karanprime]

I apologize if I am missing something fundamental. I am trying to use Jax for rigid body dynamics. As a test case, I tried manipulating the pendulum in the Brax Basics notebooks (https://colab.research.google.com/github/google/brax/blob/main/notebooks/basics.ipynb). Whatever settings I use, the triple pendulum seems to reflect from the original position (theta=0) and never crosses to the other side as expected from pendulums. I increased the number of steps (to increase the total time) and played with different initial conditions (positive and negative angles). However I cannot get it to move beyond vertical position. I have attached screenshots here:

I would really appreciate any help regarding this issue. How do I make the pendulum behave better?

[AumGupta]

Hey @karanprime,

As far as I understood, you are trying to achieve something like the image below:

Solution:

To achieve this, you just have to play with the parameter of range function used in the loop at the end of the given code.

The existing loop has 100 * substeps as parameter to range():

for i in range(100 * substeps): #Change 100 to larger values
  if i % substeps == 0:
    visualize(ax, state.x.pos, i / (100 * substeps)) #Change 100 to larger values
  state = jax.jit(pipeline.step)(pendulum, state, None)

Update it to 200 * substeps: (For Instance)

for i in range(200 * substeps): 
  if i % substeps == 0:
    visualize(ax, state.x.pos, i / (200 * substeps)) 
  state = jax.jit(pipeline.step)(pendulum, state, None)

Click to see exact code to reproduce my result

#@title { run: "auto"}

from brax.generalized import pipeline as generalized_pipeline 
from brax.positional import pipeline as positional_pipeline 
from brax.spring import pipeline as spring_pipeline 

initial_angle = 60 #@param { type:"slider", min:0, max:90, step: 1 }
pipeline = 'generalized'  #@param ["generalized", "positional", "spring"]
step_size = "5 ms" #@param [".4 ms", "1 ms", "5 ms"]

substeps = {'.4 ms': 25, '1 ms': 10, '5 ms': 2}[step_size]
stiffness = {'.4 ms': 20000, '1 ms': 10000, '5 ms': 500}[step_size]
pipeline = {'generalized': generalized_pipeline,
            'positional': positional_pipeline,
            'spring': spring_pipeline}[pipeline]

# find reasonable stiffness and damping for the given dt
link = pendulum.link.replace(
    constraint_stiffness=jp.repeat(stiffness, 3),
    constraint_vel_damping=jp.zeros((3,)))
pendulum = pendulum.replace(link=link)
pendulum = pendulum.tree_replace({'opt.timestep': 0.01 / substeps})

init_q = jp.array([initial_angle * jp.pi / 180, 0, 0])
state = jax.jit(pipeline.init)(pendulum, init_q, jp.zeros(pendulum.qd_size()))

_, ax = plt.subplots()
plt.xlim([-3, 3])
plt.ylim([-4, 0])

for i in range(200 * substeps):
  if i % substeps == 0:
    visualize(ax, state.x.pos, i / (200 * substeps))
  state = jax.jit(pipeline.step)(pendulum, state, None)

plt.title('pendulum in motion')
plt.show()

These are the exact parameters I used:

---

I hope it helped ❤️. Consider marking this as answered if it did.