Callbacks for custom behaviour during simulation

[lorycontixd]

My usage of the library (e.g. zero-order optimisation) requires the definition of a score associated with a simulation.
This method defines a score as a function that takes the simulation time and the state of the dynamics (data) at a given frame to return the score of the model for that frame. As of now the score is independent of previous values, but this can be implemented easily.

The definition of such through callbacks which are called pre, during and post simulation has the following advantages:

• Customisation of the functions that define a score, given the state of the model.
• Storage of variables during simulation
• Adaptability for different types of callbacks that a user might be interested in.

[lorycontixd]

Current usage:

def score_func(t: float, data: mujoco.MjData) -> float:
        return 0.1
    
scb = ScoreCallback(score_func) # score function = f: (t, data) -> score
tcb = TrackerCallback(["qpos"], False) # (list of tracked variables, print variables on step)

# Define simulator and set initial position
mujoco_instance = MujocoSimulator()
mujoco_instance.load_model(
    robot_model,
    s=[-0.55, 0],
    xyz_rpy=np.asarray([0, 0, initial_height, 0.5, 0, 0]),
    floor_opts={}
)
s, ds, tau = mujoco_instance.get_state()
t = mujoco_instance.get_simulation_time()
H_b = mujoco_instance.get_base()
w_b = mujoco_instance.get_base_velocity()
mujoco_instance.set_visualize_robot_flag(False)
t = 0
scb.on_simulation_start()
tcb.on_simulation_start()
while t < 4:
  # Reading robot state from simulator
  s, ds, tau = mujoco_instance.get_state()
  
  t = mujoco_instance.get_simulation_time()
  
  # Step the simulator
  mujoco_instance.step()
  scb.on_simulation_step(t, mujoco_instance.data)
  tcb.on_simulation_step(t, mujoco_instance.data)

scb.on_simulation_end()
tcb.on_simulation_end()
            

Ideally the list of callbacks should be passed to the constructor of MujocoSimulator and the callback functions should be called internally (e.g. run method)

[lorycontixd]

New EarlyStoppingCallback to finish a simulation early.

Usage:

sim = MujocoSimulator()
sim.load_model(...)

cb = EarlyStoppingCallback(lambda t, data: t > 3.0)
sim.run(6.0, callbacks=[cb])

To achieve mass callback behaviour, the main simulation loop has been grouped inside the MujocoSimulator.run() function.

[lorycontixd]

An idea for a useful callback is the ContactCallback which theoretically captures all the contacts that happened at a given frame. This should either store, alert or return all the contacts of interest during the simulation.

Input: The user passes a list of objects (str) for which he is interested of gathering contact information of (for example: ("plane", "model.right_foot" or "plane", "model.base"). If nothing is passed, all contacts can be considered by default (memory expensive) or no contacts.
Store: the class contains a list of ContactInfo for the tracked objects.
Alert: the callback raises a log, event or whatever whenever a contact between the objects has taken place, if an alert flag is active.

My doubt is whether to keep this logic inside a callback or implement it internally to the mujocoSimulator class.

Update: Could be done as both... Mujoco simulator stores it internally for eventual computations, while the callback exposes these contacts to the user in the way mentioned above.

[lorycontixd]

Currently facing a problem on TrackerCallback where tracked variables are not stored correctly. I'd like to recall that this callback takes care of storing the values of variables of interest through a simulation, in the following way:

1. Constructor: asks for a list of identifiers of the variables to be tracked during simulation
2. During simulation:
   • On start: reset the lists for storage
   • On step: retrieve the value of the variable and append it to the respective list
   • On end: -
3. Return/expose the tracked variables' values

The problem is how the values of the variables are stored. We'd obviously expect the values of a variable to be stored in the following way:

t0: [v0]
t1: [v0, v1]
...
tN: [v0, v1, ..., vN]

but for some reason, they are stored in the following way:

t0: [v0]
t1: [v1, v1]
...
tN: [vN, vN, ..., vN]

and so the returned values are all the same.

Below is an example:

Output

qpos: [0.         0.         1.49996076 1.         0.         0.
 0.        ]
qpos: [0.         0.         1.49988228 1.         0.         0.
 0.        ]
qpos: [0.         0.         1.49976456 1.         0.         0.
 0.        ]
qpos: [0.        0.        1.4996076 1.        0.        0.        0.       ]
qpos: [0.        0.        1.4994114 1.        0.        0.        0.       ]
qpos: [0.         0.         1.49917596 1.         0.         0.
 0.        ]
qpos: [0.         0.         1.49890128 1.         0.         0.
 0.        ]
qpos: [0.         0.         1.49858736 1.         0.         0.
 0.        ]
qpos: [0.        0.        1.4982342 1.        0.        0.        0.       ]
qpos: [0.        0.        1.4978418 1.        0.        0.        0.       ]
qpos: [0.         0.         1.49741016 1.         0.         0.
 0.        ]
 
Time: [0.002, 0.004, 0.006, 0.008, 0.01, 0.012, 0.014, 0.016, 0.018000000000000002, 0.020000000000000004, 0.022000000000000006]

Q position: [1.49741016 1.49741016 1.49741016 1.49741016 1.49741016 1.49741016
 1.49741016 1.49741016 1.49741016 1.49741016 1.49741016]

The lines starting with qpos are the runtime positions which are displayed correctly and different at each frame.
The final list is the history of the values, they are all the same and equal to the last value

Minimal working example

from comodo.mujocoSimulator.callbacks import ScoreCallback, TrackerCallback, EarlyStoppingCallback, ContactCallback
from comodo.robotModel.robotModel import RobotModel
from comodo.mujocoSimulator.mujocoSimulator import MujocoSimulator
import numpy as np
import pathlib
import rod
import rod.builder
import rod.builder.primitives
from rod.urdf.exporter import UrdfExporter

# Build model
dynamics = rod.Dynamics(
    damping=0.0,
    friction=0.0,
    spring_stiffness=0.0,
    spring_reference=0.0,
)
box_builder = rod.builder.primitives.BoxBuilder(
    name="root",
    mass=1.,
    x=1., y=1., z=1.,
)
box = box_builder.build_link(
    name=box_builder.name,
    pose=rod.builder.primitives.PrimitiveBuilder.build_pose(
        relative_to="__model__",
    )
).add_inertial().add_visual().add_collision().build()


rodmodel = rod.Model(
    name="model",
    canonical_link=box.name,
    link=[box],
    joint=[],
    frame=[],
)

# Save model
rod_sdf = rod.Sdf(version="1.7", model=rodmodel)
urdf_string = UrdfExporter.sdf_to_urdf_string(
    rod_sdf, pretty=True, gazebo_preserve_fixed_joints=True
)
file = pathlib.Path("robot.urdf").absolute()
with open(str(file), "w") as f:
    f.write(urdf_string)
path = pathlib.PurePosixPath(file)

tc = TrackerCallback(["qpos"], True)
escb = EarlyStoppingCallback(lambda t, iter, data, opts: iter == 10) # Stop after 10 iterations

# Load model
model = RobotModel(file, "robot", [])
mujoco_instance = MujocoSimulator()
mujoco_instance.load_model(
    model,
    s=[],
    xyz_rpy=np.array([0., 0., 1.5, 0., 0., 0.]),
    floor_opts={
        "inclination_deg" : [0, 0, 0]
    }
)
mujoco_instance.run(3.0, callbacks=[tc, escb], visualise=True)

qpos_data = tc.get_tracked_values()[1]['qpos']
t =tc.t
body_z_pos = np.array([q[2] for q in qpos_data])
print(f"Time: {t}")
print(f"Body z position: {body_z_pos}") 

[lorycontixd]

✔️ Solved!

Thanks to @fils99's enormous help, I found out that the values were stored as references inside the TrackerCallback and were modified at the beginning of each iteration when mujoco.mj_step was called. Adding a simple copy.deepcopy when appending the values in the callback class fixed the issue.