54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from typing import TYPE_CHECKING, Literal
|
||
|
|
|
||
|
|
import isaaclab.utils.math as math_utils
|
||
|
|
from isaaclab.assets import Articulation
|
||
|
|
from isaaclab.envs.mdp.events import _randomize_prop_by_op,randomize_rigid_body_com
|
||
|
|
from isaaclab.managers import SceneEntityCfg
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from isaaclab.envs import ManagerBasedEnv
|
||
|
|
|
||
|
|
|
||
|
|
def randomize_joint_default_pos(
|
||
|
|
env: ManagerBasedEnv,
|
||
|
|
env_ids: torch.Tensor | None,
|
||
|
|
asset_cfg: SceneEntityCfg,
|
||
|
|
pos_distribution_params: tuple[float, float] | None = None,
|
||
|
|
operation: Literal["add", "scale", "abs"] = "abs",
|
||
|
|
distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform",
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Randomize the joint default positions which may be different from URDF due to calibration errors.
|
||
|
|
"""
|
||
|
|
# extract the used quantities (to enable type-hinting)
|
||
|
|
asset: Articulation = env.scene[asset_cfg.name]
|
||
|
|
|
||
|
|
# save nominal value for export
|
||
|
|
asset.data.default_joint_pos_nominal = torch.clone(asset.data.default_joint_pos[0])
|
||
|
|
|
||
|
|
# resolve environment ids
|
||
|
|
if env_ids is None:
|
||
|
|
env_ids = torch.arange(env.scene.num_envs, device=asset.device)
|
||
|
|
|
||
|
|
# 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)
|
||
|
|
|
||
|
|
if pos_distribution_params is not None:
|
||
|
|
pos = asset.data.default_joint_pos.to(asset.device).clone()
|
||
|
|
pos = _randomize_prop_by_op(
|
||
|
|
pos, pos_distribution_params, env_ids, joint_ids, operation=operation, distribution=distribution
|
||
|
|
)[env_ids][:, joint_ids]
|
||
|
|
|
||
|
|
if env_ids != slice(None) and joint_ids != slice(None):
|
||
|
|
env_ids = env_ids[:, None]
|
||
|
|
asset.data.default_joint_pos[env_ids, joint_ids] = pos
|
||
|
|
# update the offset in action since it is not updated automatically
|
||
|
|
env.action_manager.get_term("joint_pos")._offset[env_ids] = pos.unsqueeze(1)
|
||
|
|
|