forked from tangger/lerobot
326aca0a48
* (unscrewing things up) (#2288) * fix: expose a function explicitly building a frame for inference * fix: first make dataset frame, then make ready for inference * fix: reducing reliance on lerobot record for policy's ouptuts too * fix: encapsulating squeezing out + device handling from predict action * fix: remove duplicated call to build_inference_frame and add a function to only perform data type handling (whole conversion is: keys matching + data type conversion) * refactor(envs): add custom-observation-size (#2167) * fix: add MockMotorBus to MockRobot * rl: first drafts * add: all components of HIL SERL * fix: actor block works * fix: less friction, less friction * add: hil-serl complete example * fix: dataset names * fix: restructuring example folder * fix: act works but found bug in how ACT works * fix: same path for both pre and postprocessors * fix: paths * add: example usage for act * add: using ACT example * fix: training examples * fix: using examples * fix: camera index * fix: rename workflows into tutorial so that the path of the files is lerobot/examples/tutorial/... * fix: upload everything in one repo * fix: model name * fix: simplify model path * add: VLAs example --------- Signed-off-by: Francesco Capuano <74058581+fracapuano@users.noreply.github.com> * fix: minor fix using named attributes * fix: change model to act * fix: named attributes for inference frame building * fix: minor fixes to smolvla * fix: small changes to pi0 * remove: old file that should have never been committed (ups sorry sorry) --------- Signed-off-by: Francesco Capuano <74058581+fracapuano@users.noreply.github.com>
58 lines
2.2 KiB
Python
58 lines
2.2 KiB
Python
import torch
|
|
|
|
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
|
|
from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
|
|
from lerobot.policies.act.modeling_act import ACTPolicy
|
|
from lerobot.policies.factory import make_pre_post_processors
|
|
from lerobot.policies.utils import build_inference_frame, make_robot_action
|
|
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
|
|
from lerobot.robots.so100_follower.so100_follower import SO100Follower
|
|
|
|
device = torch.device("mps") # or "cuda" or "cpu"
|
|
model_id = "fracapuano/robot_learning_tutorial_act"
|
|
model = ACTPolicy.from_pretrained(model_id)
|
|
|
|
dataset_id = "lerobot/svla_so101_pickplace"
|
|
# This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
|
|
dataset_metadata = LeRobotDatasetMetadata(dataset_id)
|
|
preprocess, postprocess = make_pre_post_processors(model.config, dataset_stats=dataset_metadata.stats)
|
|
|
|
# # find ports using lerobot-find-port
|
|
follower_port = ... # something like "/dev/tty.usbmodem58760431631"
|
|
|
|
# # the robot ids are used the load the right calibration files
|
|
follower_id = ... # something like "follower_so100"
|
|
|
|
MAX_EPISODES = 5
|
|
MAX_STEPS_PER_EPISODE = 20
|
|
|
|
# Robot and environment configuration
|
|
# Camera keys must match the name and resolutions of the ones used for training!
|
|
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
|
|
camera_config = {
|
|
"side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
|
|
"up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
|
|
}
|
|
|
|
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
|
|
robot = SO100Follower(robot_cfg)
|
|
robot.connect()
|
|
|
|
for _ in range(MAX_EPISODES):
|
|
for _ in range(MAX_STEPS_PER_EPISODE):
|
|
obs = robot.get_observation()
|
|
obs_frame = build_inference_frame(
|
|
observation=obs, ds_features=dataset_metadata.features, device=device
|
|
)
|
|
|
|
obs = preprocess(obs_frame)
|
|
|
|
action = model.select_action(obs)
|
|
action = postprocess(action)
|
|
|
|
action = make_robot_action(action, dataset_metadata.features)
|
|
|
|
robot.send_action(action)
|
|
|
|
print("Episode finished! Starting new episode...")
|