Adding sensors to environments

[eric-jurado]

Hello, I was using your model and was wondering how it would be possible for users to add sensors onto the fly. This is an abbreviation of the setup that I have been using, and how I currently understand how collecting sensor data works:

flight_policy_path = 'flybody-data/trained-fly-policies/flight'
wpg_pattern_path = 'flybody-data/datasets_flight-imitation/wing_pattern_fmech.npy'
n_steps = 751

 # Initialize environment
env = flight_imitation(
     None,
     wing_beat_pattern_path, 
     terminal_com_dist=float('inf')
)

# Generate a synthetic, constant speed trajectory for the fly's movement
qpos, qvel = constant_speed_trajectory(
    n_steps=n_steps+10, 
    speed=20, 
    init_pos=(0, 0, 1),
    body_rot_angle_y=-47.5, 
    control_timestep=_FLY_CONTROL_TIMESTEP)
env.task._traj_generator.set_next_trajectory(qpos,qvel)

# Wrap environment and policy
env = wrappers.SinglePrecisionWrapper(env)
env = wrappers.CanonicalSpecWrapper(env, clip=True)
policy = tf.saved_model.load(flight_policy)
policy = TestPolicyWrapper(policy)

# Initialize data and timestep variables
accel    = np.NaN*np.ones((n_steps,3))
gyro     = np.NaN*np.ones((n_steps, 3))
timestep = env.reset()

for i in tqdm(range(n_steps)):
    accel[i] = timestep.observation['walker/accelerometer']
    gyro[i]  = timestep.observation['walker/gyro']
    action = policy(timestep.observeration)
    timestep = env.step(action)

print(accel)
print(gyro)

Now, I want to directly gather data that is not normally a key within timestep.observation. My initial understanding was that I could simply inject the following into fruitfly.xml:

<sensor>
...
<force name="force_throax" site="thorax"/>
<torque name="torque_thorax" site="thorax"/>
...
</sensor>

However, this does not make timestep.observation['walker/force_thorax'] or timestep.observation['walker/torque_thorax'] exist. I've also tried to program it into the environment using,

env.task.walker.mjcf_model.sensor.add('torque', name='torque_thorax', site='thorax')

and similarly, this did not throw an error, but timestep.observation still did not change. Is there a way to make these user-added sensors observable, by somehow modifying the timestamp or otherwise?

[vaxenburg]

If you just want to gather additional data, as opposed to creating new observables to input to the policy, you don't necessarily need new sensors at all. What data are you interested in collecting?

[eric-jurado]

I would like to, for example, collect net forces and torques on various bodies of the fly. I did notice in docs/sensory-input-tracking.ipynb there is a reference to a MjData datastructure that contains simulation data that I believe would be possible to use to calculate those values manually, for each timestep, but I was wondering if there was a more convenient, built-in solution of just placing a "sensor" down on the model instead.

[vaxenburg]

My first instinct would be to use the simulation state stored in MjData, and maybe to use custom calculations or API functions like this one to compute some other quantities derived from the MjData simulation state. But maybe sometimes installing extra sensors can also be convenient?

I think you can see/access your new sensors via (though not sure the force and torque sensors measure what you're expecting):

env.physics.named.data.sensordata

The observables included in timestep.observation are in part derived from MjData.sensordata, but don't necessarily map one-to-one. See here: https://github.com/TuragaLab/flybody/blob/main/flybody/fruitfly/fruitfly.py#L581

[eric-jurado]

Thank you, this is what I was looking for. However, you mentioned that this might not give the forces and torques I am expecting; would this not work to provide, for example, an analysis of net forces and torques on the fly's center of mass, if the sensor is placed at the center of mass? Sorry if this is a basic question.

[eric-jurado]

Closing this because I figured out the answer a while back, but for anyone else who's interested:

This method will calculate net forces/torques on the fly at, or at least very near, the center of mass (the same site where body acceleration is calculated). These forces are local, though, so if you want forces in the global frame, you also have to track the body orientation env.physics.data.qpos[3:7] and rotate the forces (as vectors) by that quaternion on the same timestep.