from sim.main import InteractiveDigitalWorld
from sim.simulator import RobomimicSimulator
from sim.policy import TeleopJointPositionPolicy

if __name__ == '__main__':
    robomimic_simulator = RobomimicSimulator(
        env_name='lift'
    )
    teleop_policy = TeleopJointPositionPolicy(
        initial_position=robomimic_simulator.get_raw_state('robot0_joint_pos'),
        increment=1.0,
        keyboard_bindings=['1', '2', '3', '4', '5', '6', '7'],
        return_delta=True   # robomimic takes delta joint pos as action
    )
    playground = InteractiveDigitalWorld(
        simulator=robomimic_simulator,
        policy=teleop_policy,
        offscreen=True,
        window_size=(512, 512)
    )

    for _ in range(10):
        playground.step()

    playground.save_video(save_path='test.mp4', as_gif=False)
    playground.close()