Skip to content

Commit

Permalink
Add MIDI evaluation wrapper in example self actuated piano script.
Browse files Browse the repository at this point in the history
  • Loading branch information
kevinzakka committed Sep 7, 2023
1 parent 6aa2c51 commit eee53aa
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion examples/self_actuated_piano_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, "")
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)

0 comments on commit eee53aa

Please sign in to comment.