Views
No views yet
kosmos_ph_calvin_abcd: RoboKosMos(KosMos+Policy Head) trained on the CALVIN dataset (split ABCD).kosmos_ph_calvin_abc: RoboKosMos(KosMos+Policy Head) trained on the CALVIN dataset (split ABC).kosmos_ph_oxe-pretrain: RoboKosMos(KosMos+Policy Head) trained on the OXE-magic-soup dataset.kosmos_ph_calvin_abcd as an example:1import torch
2import json, functools
3from PIL import Image
4from robovlms.train.base_trainer import BaseTrainer
5from robovlms.data.data_utils import preprocess_image
6from robovlms.data.data_utils import get_text_function
7
8configs = josn.load(open('configs/kosmos_ph_calvin_abcd.json', 'r'))
9pretrained_path = 'checkpoints/kosmos_ph_calvin_abcd.pt'
10configs['model_load_path'] = pretrained_path
11
12model = BaseTrainer.from_checkpoint(configs)
13
14image_fn = functools.partial(
15 preprocess_image,
16 image_processor=model.model.image_processor,
17 model_type=configs["model"],
18)
19text_fn = get_text_function(model.model.tokenizer, configs["model"])
20prompt = "Task: pickup the bottle on the table"
21text_tensor, attention_mask = text_preprocess([lang])
22
23for step in range(MAX_STEPS):
24
25 image: Image.Image = get_from_side_camera(...)
26 image = image_fn([image]).unsqueeze(0)
27
28 input_dict["rgb"] = image
29 input_dict["text"] = text_tensor
30 input_dict['text_mask'] = attention_mask
31
32 ### if wrist camera is available
33 wrist_image: Image.Image = get_from_wrist_camera(...)
34 wrist_image = image_fn([wrist_image]).unsqueeze(0)
35 input_dict["hand_rgb"] = wrist_image
36
37 action = model.inference_step(input_dict)["action"]
38
39 # unormalize / reproject the action if necessary
40 from robovlms.data.data_utils import unnoramalize_action
41 if isinstance(action, tuple):
42 action = (
43 unnoramalize_action(
44 action[0], self.configs["norm_min"], self.configs["norm_max"]
45 ),
46 action[1],
47 )
48 else:
49 action = unnoramalize_action(
50 action, self.configs["norm_min"], self.configs["norm_max"]
51 )