reference/generated/torchrl.envs.MJLabWrapper

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

MJLabWrapper

torchrl.envs.MJLabWrapper(*args, **kwargs)

TorchRL wrapper for a pre-built mjlab.envs.ManagerBasedRlEnv.

Parameters:

• env – The mjlab manager-based RL environment to wrap.

• device – Torch device for actions and returned tensors. When None, it is inferred from env.device. The TorchRL device must match the mjlab simulation device; a mismatch raises an error instead of silently copying tensors across devices in the hot path.

• batch_size – Batch size of the wrapper. When None, it is inferred as [env.num_envs]. mjlab exposes a flat vectorized batch, so only one-dimensional batch sizes are accepted.

• native_autoreset – If False (default), TorchRL drives resets through the standard "_reset" mask by setting env.cfg.auto_reset to False and calling env.reset(env_ids=...) on done rows. If True, mjlab’s native auto-reset is kept and TorchRL marks the terminal "next" observation as invalid while carrying mjlab’s post-reset observation into the next root tensordict, matching the native-auto-reset contract used by Isaac Lab.

• from_pixels – If True, append RGB pixels under pixels_key at reset and step time. If the mjlab scene has an RGB CameraSensor, TorchRL uses the sensor’s batched output with shape [num_envs, H, W, 3]. Otherwise it falls back to render(), which requires num_envs == 1 and render_mode="rgb_array" because mjlab viewer rendering returns one frame for the whole scene rather than one image per environment row.

• pixels_only – If True, return only the pixel observation. Requires from_pixels=True.

• pixels_key – TensorDict key for pixel observations. Defaults to "pixels".

• pixels_sensor – Name of the mjlab RGB CameraSensor to use for from_pixels=True. If None and exactly one RGB camera sensor is present, it is selected automatically.

• allow_done_after_reset – Passed to EnvBase. Defaults to False.

• log_extras – If True, accumulate the numeric entries of the extras["log"] mapping that mjlab returns from reset and step (episode reward terms, curriculum levels, termination counts, …). Call pop_logged_extras() to retrieve their running means and clear the accumulator, typically once per training iteration. Entries are averaged over the calls in which they appeared, so metrics that mjlab emits only when an episode ends are not diluted by silent steps. Defaults to False.

mjlab reference: Zakka et al., “mjlab: A Lightweight Framework for GPU-Accelerated Robot Learning”, arXiv:2601.22074.

Examples

>>> from mjlab.envs import ManagerBasedRlEnv  
>>> from mjlab.tasks.registry import load_env_cfg  
>>> from torchrl.envs import MJLabWrapper  
>>> cfg = load_env_cfg("Mjlab-Cartpole-Balance")  
>>> cfg.scene.num_envs = 4  
>>> base_env = ManagerBasedRlEnv(cfg, device="cuda:0")  
>>> env = MJLabWrapper(base_env)  
>>> td = env.rollout(10)