From 88f7636aabde03f79a7fe71d4784b1ab64652ea4 Mon Sep 17 00:00:00 2001 From: Yu Fang Date: Wed, 16 Jul 2025 13:31:38 -0500 Subject: [PATCH] Add LocoLift task --- robosuite/environments/manipulation/lift.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/robosuite/environments/manipulation/lift.py b/robosuite/environments/manipulation/lift.py index 4ccd1f2dfd..bb22e1faac 100644 --- a/robosuite/environments/manipulation/lift.py +++ b/robosuite/environments/manipulation/lift.py @@ -1,6 +1,7 @@ from collections import OrderedDict import numpy as np +import random from robosuite.environments.manipulation.manipulation_env import ManipulationEnv from robosuite.models.arenas import TableArena @@ -438,3 +439,23 @@ def _check_success(self): # cube is higher than the table top above a margin return cube_height > table_height + 0.04 + + +class LocoLift(Lift): + + def _load_model(self): + """ + Loads an xml model, puts it in self.model + """ + self.table_full_size = random.choice([(2.0, 0.8, 0.05), (0.8, 2.0, 0.05)]) + self.table_offset = np.array( + ( + random.uniform(0, 1), + random.uniform(-0.5, 0.5), + 0.8, + ) + ) + if self.placement_initializer is not None: + self.placement_initializer.reference_pos = self.table_offset + + super()._load_model()