26 lines
709 B
Python
26 lines
709 B
Python
"""Arena composer class."""
|
|
|
|
from pathlib import Path
|
|
|
|
from dm_control import composer, mjcf
|
|
|
|
_ARENA_XML = Path(__file__).resolve().parent / "arena.xml"
|
|
|
|
|
|
class Arena(composer.Entity):
|
|
"""Base arena class."""
|
|
|
|
def _build(self, name: str = "arena") -> None:
|
|
self._mjcf_root = mjcf.from_path(str(_ARENA_XML))
|
|
if name is not None:
|
|
self._mjcf_root.model = name
|
|
|
|
def add_free_entity(self, entity) -> mjcf.Element:
|
|
"""Includes an entity as a free moving body, i.e., with a freejoint."""
|
|
frame = self.attach(entity)
|
|
frame.add("freejoint")
|
|
return frame
|
|
|
|
@property
|
|
def mjcf_model(self) -> mjcf.RootElement:
|
|
return self._mjcf_root
|