Source code for genesis_forge.mdp.observations
import torch
from genesis_forge.genesis_env import GenesisEnv
from genesis.engine.entities import RigidEntity
from genesis_forge.managers import (
PositionActionManager,
EntityManager,
ContactManager,
)
from genesis_forge.utils import entity_lin_vel, entity_ang_vel, entity_projected_gravity
"""
Entity Observations
"""
[docs]
def entity_linear_velocity(
env: GenesisEnv, entity_manager: EntityManager = None, entity_attr: str = "robot"
) -> torch.Tensor:
"""
The linear velocity of the entity's base link, in the entity's local frame.
Args:
env: The Genesis environment containing the entity
entity_manager: The entity manager for the robot/entity the observation is being computed for.
This is slightly more performant than using the `entity_attr` parameter.
entity_attr: The attribute name of the entity in the environment. This isn't necessary if `entity_manager` is provided.
Returns:
torch.Tensor: The linear velocity of the entity's base link, in the entity's local frame.
"""
if entity_manager is not None:
return entity_manager.get_linear_velocity()
entity = getattr(env, entity_attr)
return entity_lin_vel(entity)
[docs]
def entity_angular_velocity(
env: GenesisEnv, entity_manager: EntityManager = None, entity_attr: str = "robot"
) -> torch.Tensor:
"""
The angular velocity of the entity's base link, in the entity's local frame.
Args:
env: The Genesis environment containing the entity
entity_manager: The entity manager for the robot/entity the observation is being computed for.
This is slightly more performant than using the `entity_attr` parameter.
entity_attr: The attribute name of the entity in the environment. This isn't necessary if `entity_manager` is provided.
Returns:
torch.Tensor: The angular velocity of the entity's base link, in the entity's local frame.
"""
if entity_manager is not None:
return entity_manager.get_angular_velocity()
entity = getattr(env, entity_attr)
return entity_ang_vel(entity)
[docs]
def entity_projected_gravity(
env: GenesisEnv, entity_manager: EntityManager = None, entity_attr: str = "robot"
) -> torch.Tensor:
"""
The projected gravity of the entity's base link, in the entity's local frame.
Args:
env: The Genesis environment containing the entity
entity_manager: The entity manager for the robot/entity the observation is being computed for.
This is slightly more performant than using the `entity_attr` parameter.
entity_attr: The attribute name of the entity in the environment. This isn't necessary if `entity_manager` is provided.
Returns:
torch.Tensor: The projected gravity of the entity's base link, in the entity's local frame.
"""
if entity_manager is not None:
return entity_manager.get_projected_gravity()
entity = getattr(env, entity_attr)
return entity_projected_gravity(entity)
"""
DOF/Join observations
"""
[docs]
def entity_dofs_position(
env: GenesisEnv,
action_manager: PositionActionManager = None,
entity_attr: str = "robot",
dofs_idx: list[int] = None,
) -> torch.Tensor:
"""
The position of the entity's DOFs.
Args:
env: The Genesis environment containing the entity
action_manager: The action manager for the robot/entity.
This is slightly more performant than using the `entity_attr` parameter.
entity_attr: The attribute name of the entity in the environment. This isn't necessary if `action_manager` is provided.
dofs_idx: The indices of the DOFs to get the position of. This isn't necessary if `action_manager` is provided.
Returns:
torch.Tensor: The position of the entity's DOFs.
"""
if action_manager is not None:
return action_manager.get_dofs_position()
entity: RigidEntity = getattr(env, entity_attr)
return entity.get_dofs_position(dofs_idx)
[docs]
def entity_dofs_velocity(
env: GenesisEnv,
action_manager: PositionActionManager = None,
entity_attr: str = "robot",
dofs_idx: list[int] = None,
) -> torch.Tensor:
"""
The velocity of the entity's DOFs.
Args:
env: The Genesis environment containing the entity
action_manager: The action manager for the robot/entity.
This is slightly more performant than using the `entity_attr` parameter.
entity_attr: The attribute name of the entity in the environment. This isn't necessary if `action_manager` is provided.
dofs_idx: The indices of the DOFs to get the velocity of. This isn't necessary if `action_manager` is provided.
Returns:
torch.Tensor: The velocity of the entity's DOFs.
"""
if action_manager is not None:
return action_manager.get_dofs_velocity()
entity: RigidEntity = getattr(env, entity_attr)
return entity.get_dofs_velocity(dofs_idx)
[docs]
def entity_dofs_force(
env: GenesisEnv,
action_manager: PositionActionManager = None,
entity_attr: str = "robot",
dofs_idx: list[int] = None,
clip_to_max_force: bool = False,
) -> torch.Tensor:
"""
The DOF's force being experienced.
Args:
env: The Genesis environment containing the entity
action_manager: The action manager for the robot/entity.
This is slightly more performant than using the `entity_attr` parameter.
entity_attr: The attribute name of the entity in the environment. This isn't necessary if `action_manager` is provided.
dofs_idx: The indices of the DOFs to get the force of. This isn't necessary if `action_manager` is provided.
clip_to_max_force: Clip the force to the maximum force defined in the `action_manager`.
Returns:
torch.Tensor: The force of the entity's DOFs.
"""
if action_manager is not None:
return action_manager.get_dofs_force(clip_to_max_force=clip_to_max_force)
entity: RigidEntity = getattr(env, entity_attr)
return entity.get_dofs_force(dofs_idx)
"""
Actions
"""
[docs]
def current_actions(
env: GenesisEnv,
action_manager: PositionActionManager = None,
) -> torch.Tensor:
"""
The most current step actions.
"""
if action_manager is not None:
return action_manager.get_actions()
return env.actions
"""
Contacts
"""
