Feat(randomization): add articulation mass randomization functor

[yuecideng]

Description

This PR adds a randomize_articulation_mass event functor that randomizes articulation link masses within a specified range. It also adds set_mass/get_mass methods to the Articulation class for batched mass management.

Key features:

• Regex-based link selection: The link_names parameter accepts a regex pattern (str), a list of patterns (list[str]), or None (all links)
• Absolute and relative modes: Supports both absolute mass randomization and relative (additive offset from current mass)
• Batched API: New Articulation.set_mass(mass, link_names, env_ids) and Articulation.get_mass(link_names, env_ids) methods provide a clean high-level interface

Dependencies: None

Type of change

• Bug fix (non-breaking change which fixes an issue)
• Enhancement (non-breaking change which improves an existing functionality)
• New feature (non-breaking change which adds functionality)
• Breaking change (existing functionality will not work without user modification)
• Documentation update

Checklist

• I have run the black . command to format the code base.
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works
• Dependencies have been updated, if applicable.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

Adds an event functor to randomize per-link masses on articulations, along with a batched Articulation.get_mass() / Articulation.set_mass() API to support that randomization (plus docs/tests updates).

Changes:

• Add randomize_articulation_mass functor with regex-based link selection and absolute/relative modes.
• Add batched mass getters/setters on Articulation for per-link mass management.
• Extend event functor tests and documentation to cover the new mass randomization.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
tests/gym/envs/managers/test_event_functors.py	Adds mock per-link mass support and tests for randomize_articulation_mass.
embodichain/lab/sim/objects/articulation.py	Introduces Articulation.set_mass() / Articulation.get_mass() batched APIs.
embodichain/lab/sim/cfg.py	Clarifies articulation link mass/density behavior in config docstring.
embodichain/lab/gym/envs/managers/randomization/physics.py	Adds randomize_articulation_mass and related imports.
docs/source/overview/gym/event_functors.md	Documents the new event functor.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

sample_uniform is called without a device, so sampled_masses will be on CPU. In relative=True mode you then add it to current_masses returned by articulation.get_mass() (allocated on articulation.device), which will raise a device mismatch when the sim runs on CUDA. Pass an explicit device (e.g., articulation.device/env.device) to sample_uniform, and ensure both tensors are on the same device before adding.

Suggested change

	lower=mass_range[0], upper=mass_range[1], size=(num_instance, num_links)
	)
	if relative:
	# Get current mass from the articulation
	current_masses = articulation.get_mass(
	link_names=matched_link_names, env_ids=env_ids
	)
	lower=mass_range[0],
	upper=mass_range[1],
	size=(num_instance, num_links),
	device=articulation.device,
	)
	if relative:
	# Get current mass from the articulation
	current_masses = articulation.get_mass(
	link_names=matched_link_names, env_ids=env_ids
	).to(articulation.device)

[Copilot]

Robot is imported but unused in this module. Please remove it to avoid unused-import warnings and keep the imports minimal.

[Copilot]

set_mass() calls self._entities[env_idx].set_mass(name, ...), but the rest of the codebase interacts with masses through get_physical_body(...).get_mass() / RigidBody.set_mass(...). Unless the underlying articulation entity exposes set_mass, this will raise at runtime. Consider using self._entities[env_idx].get_physical_body(name).set_mass(...) for symmetry with get_mass() and with RigidObject.set_mass().

Suggested change

	self._entities[env_idx].set_mass(name, mass[i, j].item())
	self._entities[env_idx].get_physical_body(name).set_mass(
	mass[i, j].item()
	)

[Copilot]

set_mass() validates only the first dimension (env count) but not that mass is 2D and that mass.shape[1] == len(link_names). If callers pass a mismatched tensor, this will fail later with an IndexError (or silently set wrong values). Add an explicit shape check and raise via logger.log_error/ValueError before looping.

Suggested change

	if len(local_env_ids) != len(mass):
	logger.log_error(
	f"Length of env_ids {len(local_env_ids)} does not match mass length {len(mass)}."
	)
	expected_shape = (len(local_env_ids), len(link_names))
	if mass.ndim != 2 or mass.shape[0] != expected_shape[0] or mass.shape[1] != expected_shape[1]:
	msg = (
	f"Mass tensor must have shape {expected_shape}, but got {tuple(mass.shape)}."
	)
	logger.log_error(msg)
	raise ValueError(msg)

[Copilot]

set_mass() uses mass[i, j].item() inside a nested loop. On CUDA tensors this causes a device sync per element and can become very slow for many envs/links. Convert the entire tensor once (e.g., mass.cpu().numpy()), then index into that in the loop (similar to RigidObject.set_mass).

Suggested change

	for i, env_idx in enumerate(local_env_ids):
	for j, name in enumerate(link_names):
	self._entities[env_idx].set_mass(name, mass[i, j].item())
	mass_np = mass.cpu().numpy()
	for i, env_idx in enumerate(local_env_ids):
	for j, name in enumerate(link_names):
	self._entities[env_idx].set_mass(name, mass_np[i, j])

[Copilot]

get_mass() fills a tensor element-by-element in Python. If self.device is CUDA, this results in many small host->device writes. Prefer accumulating masses in a Python list / NumPy array on CPU and then torch.as_tensor(..., device=self.device) once at the end (as done in RigidObject.get_mass).

Suggested change

	mass_tensor = torch.zeros(
	(len(local_env_ids), len(link_names)),
	dtype=torch.float32,
	device=self.device,
	)
	for i, env_idx in enumerate(local_env_ids):
	for j, name in enumerate(link_names):
	mass_tensor[i, j] = (
	self._entities[env_idx].get_physical_body(name).get_mass()
	)
	return mass_tensor
	mass_values = []
	for env_idx in local_env_ids:
	env_mass_values = []
	for name in link_names:
	env_mass_values.append(
	self._entities[env_idx].get_physical_body(name).get_mass()
	)
	mass_values.append(env_mass_values)
	return torch.as_tensor(
	mass_values,
	dtype=torch.float32,
	device=self.device,
	)

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 5 comments.

Comments suppressed due to low confidence (1)

embodichain/lab/sim/objects/articulation.py:592

• default_link_masses is captured before set_dexsim_articulation_cfg() potentially overrides physical attributes (including mass). After init, default_link_masses may no longer match the articulation’s actual default masses, which can break consumers that treat it as a baseline. Consider capturing defaults after config application (or rename to reflect that it’s pre-config/USD/URDF mass).

        # Get default masses.
        self.default_link_masses = self.get_mass()

        # Determine if we should use USD properties or cfg properties.
        if not self.cfg.use_usd_properties:
            # Set articulation configuration in DexSim
            set_dexsim_articulation_cfg(entities, self.cfg)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

relative=True is documented as applying an additive offset from the current link masses, but this implementation uses articulation.default_link_masses (captured at init) as the baseline. This will give incorrect results if link masses were changed earlier (e.g., prior randomizations, asset config overrides). Use articulation.get_mass(link_names=matched_link_names, env_ids=env_ids) as the baseline when relative=True (or rename/clarify semantics if you intended “relative to defaults”).

Suggested change

	link_indices = [
	articulation.link_names.index(name) for name in matched_link_names
	]
	current_masses = articulation.default_link_masses.clone()[env_ids][
	:, link_indices
	]
	current_masses = articulation.get_mass(
	link_names=matched_link_names, env_ids=env_ids
	)

[Copilot]

set_mass() calls self._entities[env_idx].set_mass(name, ...), but other code paths manipulate link properties via get_physical_body(link_name) (and get_mass() reads via get_physical_body(name).get_mass()). If the DexSim articulation API doesn’t expose set_mass(link_name, mass) on the articulation itself, this will fail at runtime. Prefer setting mass through get_physical_body(name).set_mass(...) for symmetry with get_mass() and consistency with RigidObject.set_mass().

Suggested change

	self._entities[env_idx].set_mass(name, mass[i, j].item())
	self._entities[env_idx].get_physical_body(name).set_mass(
	mass[i, j].item()
	)

[Copilot]

Both set_mass() and get_mass() perform per-link/per-env Python loops and use .item()/scalar assignments. On CUDA, this can introduce repeated device synchronizations and many small host<->device copies. Consider converting the input tensor once to a CPU array in set_mass() (similar to RigidObject.set_mass()), and building masses on CPU then transferring once in get_mass() when device.type == "cuda".

Suggested change

	mass_tensor = torch.zeros(
	(len(local_env_ids), len(link_names)),
	dtype=torch.float32,
	device=self.device,
	)
	shape = (len(local_env_ids), len(link_names))
	if self.device.type == "cuda":
	mass_array = np.zeros(shape, dtype=np.float32)
	for i, env_idx in enumerate(local_env_ids):
	for j, name in enumerate(link_names):
	mass_array[i, j] = (
	self._entities[env_idx].get_physical_body(name).get_mass()
	)
	return torch.from_numpy(mass_array).to(device=self.device)
	mass_tensor = torch.zeros(shape, dtype=torch.float32, device=self.device)

[Copilot]

The updated docstring claims link mass defaults come from the USD/URDF unless explicitly specified, but RigidBodyAttributesCfg.mass defaults to 1.0 and set_dexsim_articulation_cfg() always applies cfg.attrs.attr() (which always sets attr.mass). Either adjust the wording to match current behavior, or change the config/API to allow “unspecified” mass so USD/URDF masses can be preserved when desired.

Suggested change

	"""Physical attributes for all links. We use default mass from the USD/URDF file if available.
	The mass and density in attrs will only be used if specified.
	"""Physical attributes for all links.
	These values are applied from ``attrs`` when configuring the articulation. In particular,
	mass and density are taken from ``RigidBodyAttributesCfg`` defaults unless explicitly changed,
	so source USD/URDF mass properties are not preserved by default through this field.

[Copilot]

Current tests for relative=True only cover the default (1.0) baseline. Add a test that first sets a non-default mass for a link (e.g., 10.0) and then applies relative=True to verify the offset is applied relative to the current mass, not an init-time cached default.

Suggested change

	def test_relative_mass_randomization_uses_current_mass(self):
	"""Test relative mass randomization uses the current mass, not a cached default."""
	env = MockEnv(num_envs=4)
	env_ids = torch.tensor([0, 1, 2, 3])
	# Set base_link mass to a known non-default value first.
	known_mass = torch.full((4, 1), 10.0)
	env.test_articulation.set_mass(
	known_mass, link_names=["base_link"], env_ids=env_ids
	)
	randomize_articulation_mass(
	env,
	env_ids,
	entity_cfg=MagicMock(uid="articulation"),
	mass_range=(-0.5, 0.5),
	link_names="base_link",
	relative=True,
	)
	# Final mass for base_link should be offset from the current mass of 10.0.
	masses = env.test_articulation.get_mass(
	link_names=["base_link"], env_ids=env_ids
	)
	assert torch.all(masses >= 9.5)
	assert torch.all(masses <= 10.5)