# Copyright (c) 2022-2024, The ORBIT Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""Common functions that can be used to enable different events.
Events include anything related to altering the simulation state. This includes changing the physics
materials, applying external forces, and resetting the state of the asset.
The functions can be passed to the :class:`omni.isaac.orbit.managers.EventTermCfg` object to enable
the event introduced by the function.
"""
from __future__ import annotations
import torch
import warnings
from typing import TYPE_CHECKING, Literal
import carb
import omni.isaac.orbit.utils.math as math_utils
from omni.isaac.orbit.actuators import ImplicitActuator
from omni.isaac.orbit.assets import Articulation, RigidObject
from omni.isaac.orbit.managers import SceneEntityCfg
from omni.isaac.orbit.terrains import TerrainImporter
if TYPE_CHECKING:
from omni.isaac.orbit.envs import BaseEnv
[docs]def randomize_rigid_body_material(
env: BaseEnv,
env_ids: torch.Tensor | None,
static_friction_range: tuple[float, float],
dynamic_friction_range: tuple[float, float],
restitution_range: tuple[float, float],
num_buckets: int,
asset_cfg: SceneEntityCfg,
):
"""Randomize the physics materials on all geometries of the asset.
This function creates a set of physics materials with random static friction, dynamic friction, and restitution
values. The number of materials is specified by ``num_buckets``. The materials are generated by sampling
uniform random values from the given ranges.
The material properties are then assigned to the geometries of the asset. The assignment is done by
creating a random integer tensor of shape (num_instances, max_num_shapes) where ``num_instances``
is the number of assets spawned and ``max_num_shapes`` is the maximum number of shapes in the asset (over
all bodies). The integer values are used as indices to select the material properties from the
material buckets.
.. attention::
This function uses CPU tensors to assign the material properties. It is recommended to use this function
only during the initialization of the environment. Otherwise, it may lead to a significant performance
overhead.
.. note::
PhysX only allows 64000 unique physics materials in the scene. If the number of materials exceeds this
limit, the simulation will crash.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# resolve environment ids
if env_ids is None:
env_ids = torch.arange(env.scene.num_envs, device="cpu")
# sample material properties from the given ranges
material_buckets = torch.zeros(num_buckets, 3)
material_buckets[:, 0].uniform_(*static_friction_range)
material_buckets[:, 1].uniform_(*dynamic_friction_range)
material_buckets[:, 2].uniform_(*restitution_range)
# create random material assignments based on the total number of shapes: num_assets x num_shapes
# note: not optimal since it creates assignments for all the shapes but only a subset is used in the body indices case.
material_ids = torch.randint(0, num_buckets, (len(env_ids), asset.root_physx_view.max_shapes))
if asset_cfg.body_ids == slice(None) or isinstance(asset, RigidObject):
# get the current materials of the bodies
materials = asset.root_physx_view.get_material_properties()
# assign the new materials
# material ids are of shape: num_env_ids x num_shapes
# material_buckets are of shape: num_buckets x 3
materials[env_ids] = material_buckets[material_ids]
# set the material properties into the physics simulation
asset.root_physx_view.set_material_properties(materials, env_ids)
elif isinstance(asset, Articulation):
# obtain number of shapes per body (needed for indexing the material properties correctly)
# note: this is a workaround since the Articulation does not provide a direct way to obtain the number of shapes
# per body. We use the physics simulation view to obtain the number of shapes per body.
num_shapes_per_body = []
for link_path in asset.root_physx_view.link_paths[0]:
link_physx_view = asset._physics_sim_view.create_rigid_body_view(link_path) # type: ignore
num_shapes_per_body.append(link_physx_view.max_shapes)
# get the current materials of the bodies
materials = asset.root_physx_view.get_material_properties()
# sample material properties from the given ranges
for body_id in asset_cfg.body_ids:
# start index of shape
start_idx = sum(num_shapes_per_body[:body_id])
# end index of shape
end_idx = start_idx + num_shapes_per_body[body_id]
# assign the new materials
# material ids are of shape: num_env_ids x num_shapes
# material_buckets are of shape: num_buckets x 3
materials[env_ids, start_idx:end_idx] = material_buckets[material_ids[:, start_idx:end_idx]]
# set the material properties into the physics simulation
asset.root_physx_view.set_material_properties(materials, env_ids)
else:
raise ValueError(
f"Randomization term 'randomize_rigid_body_material' not supported for asset: '{asset_cfg.name}'"
f" with type: '{type(asset)}' and body_ids: '{asset_cfg.body_ids}'."
)
[docs]def add_body_mass(
env: BaseEnv, env_ids: torch.Tensor | None, mass_range: tuple[float, float], asset_cfg: SceneEntityCfg
):
"""Randomize the mass of the bodies by adding a random value sampled from the given range.
.. tip::
This function uses CPU tensors to assign the body masses. It is recommended to use this function
only during the initialization of the environment.
.. deprecated:: v0.4
This function is deprecated. Please use :func:`randomize_rigid_body_mass` with ``operation="add"`` instead.
"""
msg = "Event term 'add_body_mass' is deprecated. Please use 'randomize_rigid_body_mass' with operation='add'."
warnings.warn(msg, DeprecationWarning)
carb.log_warn(msg)
# call the new function
randomize_rigid_body_mass(env, env_ids, asset_cfg, mass_range, operation="add", distribution="uniform")
[docs]def randomize_rigid_body_mass(
env: BaseEnv,
env_ids: torch.Tensor | None,
asset_cfg: SceneEntityCfg,
mass_range: tuple[float, float],
operation: Literal["add", "scale", "abs"],
distribution: Literal["uniform", "log_uniform"] = "uniform",
):
"""Randomize the mass of the bodies by adding, scaling, or setting random values.
This function allows randomizing the mass of the bodies of the asset. The function samples random values from the
given ranges and adds, scales, or sets the values into the physics simulation based on the operation.
.. tip::
This function uses CPU tensors to assign the body masses. It is recommended to use this function
only during the initialization of the environment.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# resolve body indices
if asset_cfg.body_ids == slice(None):
body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device="cpu")
else:
body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu")
# get the current masses of the bodies (num_assets, num_bodies)
masses = asset.root_physx_view.get_masses()
# sample from the given range
# note: we modify the masses in-place for all environments
# however, the setter takes care that only the masses of the specified environments are modified
masses = _randomize_prop_by_op(
masses, mass_range, env_ids, body_ids, operation=operation, distribution=distribution
)
# resolve environment ids
if env_ids is None:
env_ids = torch.arange(env.scene.num_envs, device="cpu")
# set the mass into the physics simulation
asset.root_physx_view.set_masses(masses, env_ids)
[docs]def randomize_actuator_gains(
env: BaseEnv,
env_ids: torch.Tensor | None,
asset_cfg: SceneEntityCfg,
stiffness_range: tuple[float, float] | None = None,
damping_range: tuple[float, float] | None = None,
operation: Literal["add", "scale", "abs"] = "abs",
distribution: Literal["uniform", "log_uniform"] = "uniform",
):
"""Randomize the actuator gains in an articulation by adding, scaling, or setting random values.
This function allows randomizing the actuator stiffness and damping gains.
The function samples random values from the given ranges and applies the operation to the joint properties.
It then sets the values into the actuator models. If the ranges are not provided for a particular property,
the function does not modify the property.
.. tip::
For implicit actuators, this function uses CPU tensors to assign the actuator gains into the simulation.
In such cases, it is recommended to use this function only during the initialization of the environment.
Raises:
NotImplementedError: If the joint indices are in explicit motor mode. This operation is currently
not supported for explicit actuator models.
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
# resolve joint indices
if asset_cfg.joint_ids == slice(None):
joint_ids_list = range(asset.num_joints)
joint_ids = slice(None) # for optimization purposes
else:
joint_ids_list = asset_cfg.joint_ids
joint_ids = torch.tensor(asset_cfg.joint_ids, dtype=torch.int, device=asset.device)
# check if none of the joint indices are in explicit motor mode
for joint_index in joint_ids_list:
for act_name, actuator in asset.actuators.items():
# if joint indices are a slice (i.e., all joints are captured) or the joint index is in the actuator
if actuator.joint_indices == slice(None) or joint_index in actuator.joint_indices:
if not isinstance(actuator, ImplicitActuator):
raise NotImplementedError(
"Event term 'randomize_actuator_stiffness_and_damping' is performed on asset"
f" '{asset_cfg.name}' on the joint '{asset.joint_names[joint_index]}' ('{joint_index}') which"
f" uses an explicit actuator model '{act_name}<{actuator.__class__.__name__}>'. This operation"
" is currently not supported for explicit actuator models."
)
# sample joint properties from the given ranges and set into the physics simulation
# -- stiffness
if stiffness_range is not None:
stiffness = asset.root_physx_view.get_dof_stiffnesses().to(asset.device)
stiffness = _randomize_prop_by_op(
stiffness, stiffness_range, env_ids, joint_ids, operation=operation, distribution=distribution
)
asset.write_joint_stiffness_to_sim(stiffness, joint_ids=joint_ids, env_ids=env_ids)
# -- damping
if damping_range is not None:
damping = asset.root_physx_view.get_dof_dampings().to(asset.device)
damping = _randomize_prop_by_op(
damping, damping_range, env_ids, joint_ids, operation=operation, distribution=distribution
)
asset.write_joint_damping_to_sim(damping, joint_ids=joint_ids, env_ids=env_ids)
[docs]def randomize_joint_parameters(
env: BaseEnv,
env_ids: torch.Tensor | None,
asset_cfg: SceneEntityCfg,
friction_range: tuple[float, float] | None = None,
armature_range: tuple[float, float] | None = None,
operation: Literal["add", "scale", "abs"] = "abs",
distribution: Literal["uniform", "log_uniform"] = "uniform",
):
"""Randomize the joint parameters of an articulation by adding, scaling, or setting random values.
This function allows randomizing the joint parameters (friction and armature) of the asset. These correspond
to the physics engine joint properties that affect the joint behavior.
The function samples random values from the given ranges and applies the operation to the joint properties.
It then sets the values into the physics simulation. If the ranges are not provided for a
particular property, the function does not modify the property.
.. tip::
This function uses CPU tensors to assign the joint properties. It is recommended to use this function
only during the initialization of the environment.
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
# resolve joint indices
if asset_cfg.joint_ids == slice(None):
joint_ids = slice(None) # for optimization purposes
else:
joint_ids = torch.tensor(asset_cfg.joint_ids, dtype=torch.int, device=asset.device)
# sample joint properties from the given ranges and set into the physics simulation
# -- friction
if friction_range is not None:
friction = asset.root_physx_view.get_dof_friction_coefficients().to(asset.device)
friction = _randomize_prop_by_op(
friction, friction_range, env_ids, joint_ids, operation=operation, distribution=distribution
)
asset.write_joint_friction_to_sim(friction, joint_ids=joint_ids, env_ids=env_ids)
# -- armature
if armature_range is not None:
armature = asset.root_physx_view.get_dof_armatures().to(asset.device)
armature = _randomize_prop_by_op(
armature, armature_range, env_ids, joint_ids, operation=operation, distribution=distribution
)
asset.write_joint_armature_to_sim(armature, joint_ids=joint_ids, env_ids=env_ids)
[docs]def apply_external_force_torque(
env: BaseEnv,
env_ids: torch.Tensor,
force_range: tuple[float, float],
torque_range: tuple[float, float],
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
):
"""Randomize the external forces and torques applied to the bodies.
This function creates a set of random forces and torques sampled from the given ranges. The number of forces
and torques is equal to the number of bodies times the number of environments. The forces and torques are
applied to the bodies by calling ``asset.set_external_force_and_torque``. The forces and torques are only
applied when ``asset.write_data_to_sim()`` is called in the environment.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# resolve environment ids
if env_ids is None:
env_ids = torch.arange(env.scene.num_envs, device=asset.device)
# resolve number of bodies
num_bodies = len(asset_cfg.body_ids) if isinstance(asset_cfg.body_ids, list) else asset.num_bodies
# sample random forces and torques
size = (len(env_ids), num_bodies, 3)
forces = math_utils.sample_uniform(*force_range, size, asset.device)
torques = math_utils.sample_uniform(*torque_range, size, asset.device)
# set the forces and torques into the buffers
# note: these are only applied when you call: `asset.write_data_to_sim()`
asset.set_external_force_and_torque(forces, torques, env_ids=env_ids, body_ids=asset_cfg.body_ids)
[docs]def push_by_setting_velocity(
env: BaseEnv,
env_ids: torch.Tensor,
velocity_range: dict[str, tuple[float, float]],
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
):
"""Push the asset by setting the root velocity to a random value within the given ranges.
This creates an effect similar to pushing the asset with a random impulse that changes the asset's velocity.
It samples the root velocity from the given ranges and sets the velocity into the physics simulation.
The function takes a dictionary of velocity ranges for each axis and rotation. The keys of the dictionary
are ``x``, ``y``, ``z``, ``roll``, ``pitch``, and ``yaw``. The values are tuples of the form ``(min, max)``.
If the dictionary does not contain a key, the velocity is set to zero for that axis.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# velocities
vel_w = asset.data.root_vel_w[env_ids]
# sample random velocities
range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
ranges = torch.tensor(range_list, device=asset.device)
vel_w[:] = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], vel_w.shape, device=asset.device)
# set the velocities into the physics simulation
asset.write_root_velocity_to_sim(vel_w, env_ids=env_ids)
[docs]def reset_root_state_with_random_orientation(
env: BaseEnv,
env_ids: torch.Tensor,
pose_range: dict[str, tuple[float, float]],
velocity_range: dict[str, tuple[float, float]],
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
):
"""Reset the asset root position and velocities sampled randomly within the given ranges
and the asset root orientation sampled randomly from the SO(3).
This function randomizes the root position and velocity of the asset.
* It samples the root position from the given ranges and adds them to the default root position, before setting
them into the physics simulation.
* It samples the root orientation uniformly from the SO(3) and sets them into the physics simulation.
* It samples the root velocity from the given ranges and sets them into the physics simulation.
The function takes a dictionary of position and velocity ranges for each axis and rotation:
* :attr:`pose_range` - a dictionary of position ranges for each axis. The keys of the dictionary are ``x``,
``y``, and ``z``. The orientation is sampled uniformly from the SO(3).
* :attr:`velocity_range` - a dictionary of velocity ranges for each axis and rotation. The keys of the dictionary
are ``x``, ``y``, ``z``, ``roll``, ``pitch``, and ``yaw``.
The values are tuples of the form ``(min, max)``. If the dictionary does not contain a particular key,
the position is set to zero for that axis.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# get default root state
root_states = asset.data.default_root_state[env_ids].clone()
# poses
range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]]
ranges = torch.tensor(range_list, device=asset.device)
rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 3), device=asset.device)
positions = root_states[:, 0:3] + env.scene.env_origins[env_ids] + rand_samples
orientations = math_utils.random_orientation(len(env_ids), device=asset.device)
# velocities
range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
ranges = torch.tensor(range_list, device=asset.device)
rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=asset.device)
velocities = root_states[:, 7:13] + rand_samples
# set into the physics simulation
asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids)
asset.write_root_velocity_to_sim(velocities, env_ids=env_ids)
[docs]def reset_root_state_from_terrain(
env: BaseEnv,
env_ids: torch.Tensor,
pose_range: dict[str, tuple[float, float]],
velocity_range: dict[str, tuple[float, float]],
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
):
"""Reset the asset root state by sampling a random valid pose from the terrain.
This function samples a random valid pose(based on flat patches) from the terrain and sets the root state
of the asset to this position. The function also samples random velocities from the given ranges and sets them
into the physics simulation.
The function takes a dictionary of position and velocity ranges for each axis and rotation:
* :attr:`pose_range` - a dictionary of pose ranges for each axis. The keys of the dictionary are ``roll``,
``pitch``, and ``yaw``. The position is sampled from the flat patches of the terrain.
* :attr:`velocity_range` - a dictionary of velocity ranges for each axis and rotation. The keys of the dictionary
are ``x``, ``y``, ``z``, ``roll``, ``pitch``, and ``yaw``.
The values are tuples of the form ``(min, max)``. If the dictionary does not contain a particular key,
the position is set to zero for that axis.
Note:
The function expects the terrain to have valid flat patches under the key "init_pos". The flat patches
are used to sample the random pose for the robot.
Raises:
ValueError: If the terrain does not have valid flat patches under the key "init_pos".
"""
# access the used quantities (to enable type-hinting)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
terrain: TerrainImporter = env.scene.terrain
# obtain all flat patches corresponding to the valid poses
valid_positions: torch.Tensor = terrain.flat_patches.get("init_pos")
if valid_positions is None:
raise ValueError(
"The event term 'reset_root_state_from_terrain' requires valid flat patches under 'init_pos'."
f" Found: {list(terrain.flat_patches.keys())}"
)
# sample random valid poses
ids = torch.randint(0, valid_positions.shape[2], size=(len(env_ids),), device=env.device)
positions = valid_positions[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids], ids]
positions += asset.data.default_root_state[env_ids, :3]
# sample random orientations
range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["roll", "pitch", "yaw"]]
ranges = torch.tensor(range_list, device=asset.device)
rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 3), device=asset.device)
# convert to quaternions
orientations = math_utils.quat_from_euler_xyz(rand_samples[:, 0], rand_samples[:, 1], rand_samples[:, 2])
# sample random velocities
range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
ranges = torch.tensor(range_list, device=asset.device)
rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=asset.device)
velocities = asset.data.default_root_state[:, 7:13] + rand_samples
# set into the physics simulation
asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids)
asset.write_root_velocity_to_sim(velocities, env_ids=env_ids)
[docs]def reset_joints_by_scale(
env: BaseEnv,
env_ids: torch.Tensor,
position_range: tuple[float, float],
velocity_range: tuple[float, float],
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
):
"""Reset the robot joints by scaling the default position and velocity by the given ranges.
This function samples random values from the given ranges and scales the default joint positions and velocities
by these values. The scaled values are then set into the physics simulation.
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
# get default joint state
joint_pos = asset.data.default_joint_pos[env_ids].clone()
joint_vel = asset.data.default_joint_vel[env_ids].clone()
# scale these values randomly
joint_pos *= math_utils.sample_uniform(*position_range, joint_pos.shape, joint_pos.device)
joint_vel *= math_utils.sample_uniform(*velocity_range, joint_vel.shape, joint_vel.device)
# clamp joint pos to limits
joint_pos_limits = asset.data.soft_joint_pos_limits[env_ids]
joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1])
# clamp joint vel to limits
joint_vel_limits = asset.data.soft_joint_vel_limits[env_ids]
joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits)
# set into the physics simulation
asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids)
[docs]def reset_joints_by_offset(
env: BaseEnv,
env_ids: torch.Tensor,
position_range: tuple[float, float],
velocity_range: tuple[float, float],
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
):
"""Reset the robot joints with offsets around the default position and velocity by the given ranges.
This function samples random values from the given ranges and biases the default joint positions and velocities
by these values. The biased values are then set into the physics simulation.
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
# get default joint state
joint_pos = asset.data.default_joint_pos[env_ids].clone()
joint_vel = asset.data.default_joint_vel[env_ids].clone()
# bias these values randomly
joint_pos += math_utils.sample_uniform(*position_range, joint_pos.shape, joint_pos.device)
joint_vel += math_utils.sample_uniform(*velocity_range, joint_vel.shape, joint_vel.device)
# clamp joint pos to limits
joint_pos_limits = asset.data.soft_joint_pos_limits[env_ids]
joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1])
# clamp joint vel to limits
joint_vel_limits = asset.data.soft_joint_vel_limits[env_ids]
joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits)
# set into the physics simulation
asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids)
[docs]def reset_scene_to_default(env: BaseEnv, env_ids: torch.Tensor):
"""Reset the scene to the default state specified in the scene configuration."""
# rigid bodies
for rigid_object in env.scene.rigid_objects.values():
# obtain default and deal with the offset for env origins
default_root_state = rigid_object.data.default_root_state[env_ids].clone()
default_root_state[:, 0:3] += env.scene.env_origins[env_ids]
# set into the physics simulation
rigid_object.write_root_state_to_sim(default_root_state, env_ids=env_ids)
# articulations
for articulation_asset in env.scene.articulations.values():
# obtain default and deal with the offset for env origins
default_root_state = articulation_asset.data.default_root_state[env_ids].clone()
default_root_state[:, 0:3] += env.scene.env_origins[env_ids]
# set into the physics simulation
articulation_asset.write_root_state_to_sim(default_root_state, env_ids=env_ids)
# obtain default joint positions
default_joint_pos = articulation_asset.data.default_joint_pos[env_ids].clone()
default_joint_vel = articulation_asset.data.default_joint_vel[env_ids].clone()
# set into the physics simulation
articulation_asset.write_joint_state_to_sim(default_joint_pos, default_joint_vel, env_ids=env_ids)
"""
Internal helper functions.
"""
def _randomize_prop_by_op(
data: torch.Tensor,
sample_range: tuple[float, float],
dim_0_ids: torch.Tensor | None,
dim_1_ids: torch.Tensor | slice,
operation: Literal["add", "scale", "abs"],
distribution: Literal["uniform", "log_uniform"],
) -> torch.Tensor:
"""Perform data randomization based on the given operation and distribution.
Args:
data: The data tensor to be randomized. Shape is (dim_0, dim_1).
sample_range: The range to sample the random values from.
dim_0_ids: The indices of the first dimension to randomize.
dim_1_ids: The indices of the second dimension to randomize.
operation: The operation to perform on the data. Options: 'add', 'scale', 'abs'.
distribution: The distribution to sample the random values from. Options: 'uniform', 'log_uniform'.
Returns:
The data tensor after randomization. Shape is (dim_0, dim_1).
Raises:
NotImplementedError: If the operation or distribution is not supported.
"""
# resolve shape
# -- dim 0
if dim_0_ids is None:
n_dim_0 = data.shape[0]
dim_0_ids = slice(None)
else:
n_dim_0 = len(dim_0_ids)
# -- dim 1
if isinstance(dim_1_ids, slice):
n_dim_1 = data.shape[1]
else:
n_dim_1 = len(dim_1_ids)
# resolve the distribution
if distribution == "uniform":
dist_fn = math_utils.sample_uniform
elif distribution == "log_uniform":
dist_fn = math_utils.sample_log_uniform
else:
raise NotImplementedError(
f"Unknown distribution: '{distribution}' for joint properties randomization."
" Please use 'uniform' or 'log_uniform'."
)
# perform the operation
if operation == "add":
data[dim_0_ids, dim_1_ids] += dist_fn(*sample_range, (n_dim_0, n_dim_1), device=data.device)
elif operation == "scale":
data[dim_0_ids, dim_1_ids] *= dist_fn(*sample_range, (n_dim_0, n_dim_1), device=data.device)
elif operation == "abs":
data[dim_0_ids, dim_1_ids] = dist_fn(*sample_range, (n_dim_0, n_dim_1), device=data.device)
else:
raise NotImplementedError(
f"Unknown operation: '{operation}' for property randomization. Please use 'add', 'scale', or 'abs'."
)
return data
