Fix test play

minestudio/tutorials/simulator/test_play.py

@@ -1,53 +1,30 @@
 import numpy as np
 from minestudio.simulator import MinecraftSim
 from minestudio.simulator.callbacks import (
-    PlayCallback, RecordCallback, PointCallback, SpeedTestCallback, SummonMobsCallback, MaskActionsCallback, RewardsCallback, CommandsCallback, TaskCallback, JudgeResetCallback, FastResetCallback
+    PlayCallback, RecordCallback, PointCallback, PlaySegmentCallback
 )
 from minestudio.simulator.utils.gui import RecordDrawCall, CommandModeDrawCall, SegmentDrawCall
 from functools import partial
-from minestudio.models import load_vpt_policy, load_rocket_policy
+from minestudio.models import load_openai_policy, load_rocket_policy
 if __name__ == '__main__':
     agent_generator = partial(
         load_rocket_policy,
-        ckpt_path="/home/zhwang/Desktop/zhancun/jarvisbase/pretrained/rocket_12-01.ckpt",
+        ckpt_path = 'YOUR CKPT PATH',
     )
     sim = MinecraftSim(
-            obs_size=(128, 128), 
-            preferred_spawn_biome="plains", 
-            action_type = "agent",
-            timestep_limit=1000,
-            callbacks=[
-                #SpeedTestCallback(50), 
-                SummonMobsCallback([{'name': 'cow', 'number': 10, 'range_x': [-5, 5], 'range_z': [-5, 5]}]),
-                MaskActionsCallback(inventory=0), 
-                RewardsCallback([{
-                    'event': 'kill_entity', 
-                    'objects': ['cow', 'sheep'], 
-                    'reward': 5.0, 
-                    'identity': 'kill sheep or cow', 
-                    'max_reward_times': 5, 
-                }]),
-                CommandsCallback(commands=[
-                    '/give @p minecraft:iron_sword 1',
-                    '/give @p minecraft:diamond 64',
-                ]), 
-                FastResetCallback(
-                    biomes=['mountains'],
-                    random_tp_range=1000,
-                ), 
-                JudgeResetCallback(600),
-            ]
-        )
-    sim.reset()
-    for i in range(100):
-        for j in range(100):
-            obs, reward, terminated, truncated, info = sim.step(sim.noop_action())
-        obs, info = sim.reset()
-        print("info['kill_entity']: ", info['kill_entity'])
+        obs_size=(224, 224),
+        action_type="env",
+        callbacks=[
+            PlaySegmentCallback(sam_path='YOUR SAM PATH', sam_choice='small'),
+            PlayCallback(agent_generator=agent_generator, extra_draw_call=[RecordDrawCall, CommandModeDrawCall, SegmentDrawCall]),
+            RecordCallback(record_path='./output', recording=False),
+        ]
+    )
+    obs, info = sim.reset()
     terminated = False
 
-    # while not terminated:
-    #     action = None
-    #     obs, reward, terminated, truncated, info = sim.step(action)
+    while not terminated:
+        action = None
+        obs, reward, terminated, truncated, info = sim.step(action)
 
     sim.close()