Add MIDI evaluation wrapper in example self actuated piano script.

kevinzakka · kevinzakka · commit eee53aa07c40 · 2023-09-06T20:08:31.000-07:00
diff --git a/examples/self_actuated_piano_env.py b/examples/self_actuated_piano_env.py
@@ -23,7 +23,7 @@
 
 from robopianist import music, viewer
 from robopianist.suite.tasks import self_actuated_piano
-from robopianist.wrappers import PianoSoundVideoWrapper
+from robopianist.wrappers import MidiEvaluationWrapper, PianoSoundVideoWrapper
 
 _FILE = flags.DEFINE_string("file", "TwinkleTwinkleRousseau", "")
 _RECORD = flags.DEFINE_bool("record", False, "")
@@ -54,6 +54,7 @@ def main(_) -> None:
     env = composer_utils.Environment(
         recompile_physics=False, task=task, strip_singleton_obs_buffer_dim=True
     )
+    env = MidiEvaluationWrapper(env)
     if _RECORD.value:
         env = PianoSoundVideoWrapper(
             env,
@@ -104,6 +105,9 @@ def __call__(self, timestep: dm_env.TimeStep) -> np.ndarray:
             action = policy(timestep)
             timestep = env.step(action)
 
+    for k, v in env.get_musical_metrics().items():
+        np.testing.assert_equal(v, 1.0)
+
 
 if __name__ == "__main__":
     app.run(main)
