brax env and state as a member of a class/pytree. Are they static?

[yuichiro-student]

I would like to have a class that has a method that takes in generalized position, velocity, and action as its input and outputs the generalized position, velocity, and rotation at the next time step.

I have a code as below. I am carrying brax env and state as static (not dynamic arrays registered as children).
Is this correct from the developer's perspective? Because it seems that they could be registered as dynamic arrays (children).

Thank you.

class Ant:
    def __init__(self, state, env,  n_x, n_u):
        self.env = env
        self.state = state
        self.n_x = n_x
        self.n_u = n_u

    def dynamics(self,x,u):
        q = x[:15]
        qd = x[15:self.n_x] 

        state = self.state.tree_replace({'pipeline_state.qpos': q, 'pipeline_state.qvel': qd})
        state = self.env.step(state,u)
        
        # extract gen. pos. and velocity.
        q_new = state.pipeline_state.qpos
        qd_new = state.pipeline_state.qvel
        x_pos = state.pipeline_state.x.pos.flatten()
        x_rot = state.pipeline_state.x.rot.flatten()
         
        x_new = jnp.concatenate((q_new, qd_new, x_pos, x_rot))
        return x_new
  

    def _tree_flatten(self):
        children = ()
        aux_data = {'state':self.state, 'env':self.env, 'n_x':self.n_x, 'n_u':self.n_u}
        return (children, aux_data)

    @classmethod
    def _tree_unflatten(cls, aux_data, children):
        return cls(*children, **aux_data)
        
tree_util.register_pytree_node(Ant,
                               Ant._tree_flatten,
                               Ant._tree_unflatten