Building on BiGym 2.0#
Your own tasks and lower-body controllers can live in a separate package
that depends on bigym, with no change to BiGym 2.0. The code samples below come
from tests/fixtures/external_extension.py,
which the test suite builds and steps.
Structure#
my_lab/
pyproject.toml # depends on bigym
src/my_lab/
__init__.py
tasks.py # TaskSpecs and their env classes
wbc.py # lower-body controller and its BackendBinding
Depend on bigym#
uv init --package my_lab
cd my_lab
uv add "bigym @ git+https://github.com/swirl-uk/BiGym2"
git clone https://github.com/swirl-uk/BiGym2.git
uv add --editable /path/to/BiGym2
Extras go in the requirement, e.g.
"bigym[vr] @ git+https://github.com/swirl-uk/BiGym2" for bigym-collect
(see Installation). Keep the [build-system]
table that uv init --package writes. Without it your package is not
installed into the environment, and the BiGym 2.0 command-line tools cannot
import my_lab.
Configure an existing task#
Keyword overrides or an EnvConfig change any
setting of a task (see Overrides):
from bigym.loco import make
env = make("move_plate", camera_keys=("head",), controller={"cmd_clip": 0.5})
Add a task#
A task is a TaskSpec: a BiGymEnv subclass, an
episode budget in env steps, and the EnvConfig fields where the task
departs from the defaults (in make’s override form). The shortest route is
to subclass an existing task and change its class attributes or its
_initialize_env, _on_reset, _success and _fail hooks
(bigym/envs/move_plates.py
shows all four):
class PreciseMovePlateG1(MovePlateG1):
"""move_plate with half the distance tolerance at the target rack."""
_SUCCESSFUL_DIST = 0.025
TASK = TaskSpec(
PreciseMovePlateG1,
episode_length=17000,
overrides={"controller": {"pitch_command": False}},
)
_initialize_env adds props to the scene’s mujoco.MjSpec (self.spec),
which is compiled once afterwards. The other hooks read and write the
compiled arrays through self.model.bind(element) and
self.data.bind(element). Element names are full scene names, such as
"dishwasher/door". After writing state, call self.simulation.forward()
before reading derived quantities such as xpos.
Build it by a registered name or by its "pkg.module:ATTR" reference:
from bigym.loco import make, register_task
from my_lab.tasks import TASK
register_task("precise_move_plate", TASK)
env = make("precise_move_plate")
env = make("my_lab.tasks:TASK") # imported on first use, no registration
A registered name exists only in processes that ran register_task. The
reference works in any process with my_lab installed, including the ones
the BiGym 2.0 tools start. It is also the name the env records in its metadata,
so a recorded batch rebuilds the same task. Use the reference on the
command line.
A TaskSpec can also pin the task’s controller:
overrides={"controller": {"backend": "my_lab.wbc:HOLD_POSE"}}.
Add a lower-body controller#
A controller subclasses LowerBodyBase (the
contract is in Adding a backend). This one
holds the legs and waist at a standing pose:
# Hip pitch, hip roll, hip yaw, knee, ankle pitch, ankle roll (left, then
# right), then waist yaw, roll, pitch: the order of G1_GROOT_WBC_JOINT_NAMES.
STANDING_POSE = np.array(
[-0.1, 0.0, 0.0, 0.3, -0.2, 0.0] * 2 + [0.0, 0.0, 0.0], dtype=np.float32
)
class HoldPoseController(LowerBodyBase):
"""Holds the legs and waist at a standing pose and ignores the command."""
STATEFUL = {"cmd": "command", "height_cmd": "height_command"}
def __init__(self, env, *, control_dt):
joints = G1_GROOT_WBC_JOINT_NAMES
self._env = env
self.control_dt = control_dt
self.controlled_joints = joints
self.command_spec = velocity_spec(
vx=(-1.0, 1.0), vy=(-1.0, 1.0), wz=(-1.0, 1.0), height=(0.4, 1.0)
)
self.velocity_clip = self.yaw_rate_clip = 1.0
self.qpos_addresses, self.dof_addresses = self.build_joint_addresses(joints)
low, high = self.build_joint_ranges(joints)
self.controlled_range_low, self.controlled_range_high = low, high
_, self.base_dof_addresses = self.build_joint_addresses(
("pelvis_x", "pelvis_y", "pelvis_z")
)
orientation = self.find_sensor("orientation", dim=4)
gyro = self.find_sensor("angular-velocity", dim=3)
assert orientation is not None and gyro is not None, "no pelvis IMU"
self.orientation_sensor_address, self.gyro_sensor_address = orientation, gyro
self.last_action = np.zeros(len(STANDING_POSE), dtype=np.float32)
self.reset()
def reset(self):
"""Clear the command and put the joints in the standing pose."""
self.command = np.zeros(3, dtype=np.float32)
self.height_command = 0.74
self.apply_pose(
self.qpos_addresses,
self.dof_addresses,
STANDING_POSE,
G1_GROOT_WBC_JOINT_NAMES,
)
def step(self):
"""The standing pose, in controlled_joints order."""
return STANDING_POSE.copy()
def get_base_obs(self):
"""Pelvis linear velocity, angular velocity and gravity, pelvis frame."""
data = self._env.data
quat = data.sensordata[
self.orientation_sensor_address : self.orientation_sensor_address + 4
]
lin_vel = quat_rotate_inverse_wxyz(quat, data.qvel[self.base_dof_addresses])
ang_vel = data.sensordata[
self.gyro_sensor_address : self.gyro_sensor_address + 3
]
gravity = quat_rotate_inverse_wxyz(quat, np.array([0.0, 0.0, -1.0]))
return lin_vel, ang_vel.astype(np.float32), gravity
A learned controller loads its policy in __init__, runs it in step(),
and returns its weight files from weight_files(). The substrate
fingerprint records their SHA-256.
bigym/loco/adapters/groot_wbc.py
does all three.
A BackendBinding puts the controller in the env: the
robot models it drives, the robot class it needs (which joints are
actuated, with which PD gains) and how to build it:
class HoldPoseBinding(BackendBinding):
"""HoldPoseController on the G1 Dex1, legs and waist on the GR00T-WBC gains."""
robot_models = ("g1_dex1",)
def robot_cls(self, robot_cls, config):
"""The G1 Dex1 with the controlled joints actuated."""
assert issubclass(robot_cls, G1Dex1)
return robot_cls.variant(
actuated=G1_GROOT_WBC_JOINT_NAMES, joint_pd=G1_GROOT_WBC_JOINT_PD
)
def build_controller(self, env, config, *, control_dt):
"""A HoldPoseController stepping at the env's control rate."""
return HoldPoseController(env, control_dt=control_dt)
HOLD_POSE = HoldPoseBinding()
Select it like a task, by reference or by a name given to
register_backend():
from bigym.loco import make, register_backend
from my_lab.wbc import HOLD_POSE
env = make("move_plate", controller={"backend": "my_lab.wbc:HOLD_POSE"})
register_backend("hold_pose", HOLD_POSE)
env = make("move_plate", controller={"backend": "hold_pose"})
Use them with the BiGym 2.0 tools#
Tool |
Your task |
Your controller |
|---|---|---|
|
name or reference |
name or reference |
|
|
|
|
|
|
|
read from the batch metadata |
read from the batch metadata |
|
from a dataset repo you publish (below) |
n/a |
|
built-in tasks only |
|
bigym-agent needs a one-sentence brief and published demonstrations for
each task it runs, and both exist for the built-in tasks only.
The evaluation runner loads the policy from a module:factory import spec
whose factory returns a callable from timestep to action:
uv run python -m bigym.loco.eval.runner --task my_lab.tasks:TASK \
--method my-method --policy my_lab.policies:make_policy --out result.json
Collecting demonstrations for your task#
On a machine set up for VR collection, with
bigym[vr] in your project:
uv run bigym-collect --task my_lab.tasks:TASK --export-lerobot \
--task-text "Move the plate into the other rack."
The batch goes to ./bigym_demos/my_lab.tasks-TASK/<timestamp>. File and
directory names replace the colon with a dash, and the batch metadata keeps
the exact task name. After the session the collector cuts a training view
to the task’s success hold (<timestamp>_hold1s for a 1 s hold) and
exports it to LeRobot in <timestamp>_hold1s_lerobot. --task-text is the
language instruction written to the LeRobot dataset. The built-in instruction
table covers only the built-in tasks.
env.get_demos() reads demonstrations from a Hugging Face dataset repo with
one folder per task. Upload the export as the folder my_lab.tasks-TASK/
of a dataset repo of yours and set BIGYM_DATASET_REPO to it:
BIGYM_DATASET_REPO=my-org/my-lab-demos uv run python train.py