reaf-sim 0.0.1__py3-none-any.whl

reaf/simulation/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Robotics Environment Authoring Framework - Simulation Modules."""

reaf/simulation/torch/__init__.py

@@ -0,0 +1,153 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""reaf-sim vectorized simulation framework."""
+
+from reaf.simulation.torch._src.builders import flat_scene
+from reaf.simulation.torch._src.common import action_history
+from reaf.simulation.torch._src.common import actuator_device
+from reaf.simulation.torch._src.common import articulated_device
+from reaf.simulation.torch._src.common import batched_camera_device
+from reaf.simulation.torch._src.common import camera_device
+from reaf.simulation.torch._src.common import clock_device
+from reaf.simulation.torch._src.common import contact_sensor_device
+from reaf.simulation.torch._src.common import default_discount_provider
+from reaf.simulation.torch._src.common import delay_buffer
+from reaf.simulation.torch._src.common import free_device
+from reaf.simulation.torch._src.common import functional_device
+from reaf.simulation.torch._src.common import functional_observation
+from reaf.simulation.torch._src.common import functional_reset
+from reaf.simulation.torch._src.common import functional_reward
+from reaf.simulation.torch._src.common import functional_termination
+from reaf.simulation.torch._src.common import inside_site_device
+from reaf.simulation.torch._src.common import math
+from reaf.simulation.torch._src.common import model_randomization
+from reaf.simulation.torch._src.common import mujoco_camera_device
+from reaf.simulation.torch._src.common import multi_asset_device
+from reaf.simulation.torch._src.common import noise
+from reaf.simulation.torch._src.common import processors
+from reaf.simulation.torch._src.common import resetter
+from reaf.simulation.torch._src.common import site_device
+from reaf.simulation.torch._src.common import spec_helper
+from reaf.simulation.torch._src.common import tensors
+from reaf.simulation.torch._src.common import timeout_termination_checker
+from reaf.simulation.torch._src.common import zero_reward_provider
+from reaf.simulation.torch._src.common.action_history import ActionHistory
+from reaf.simulation.torch._src.common.actuator_device import ActuatorDevice
+from reaf.simulation.torch._src.common.actuator_device import ControlMode
+from reaf.simulation.torch._src.common.actuator_device import TransmissionType
+from reaf.simulation.torch._src.common.articulated_device import ArticulatedDevice
+from reaf.simulation.torch._src.common.batched_camera_device import BatchedMujocoCameraDevice
+from reaf.simulation.torch._src.common.camera_device import CameraDevice
+from reaf.simulation.torch._src.common.clock_device import ClockDevice
+from reaf.simulation.torch._src.common.contact_sensor_device import ContactField
+from reaf.simulation.torch._src.common.contact_sensor_device import ContactSensorDevice
+from reaf.simulation.torch._src.common.contact_sensor_device import ObjType
+from reaf.simulation.torch._src.common.contact_sensor_device import ReduceMode
+from reaf.simulation.torch._src.common.default_discount_provider import DefaultDiscountProvider
+from reaf.simulation.torch._src.common.delay_buffer import DelayBuffer
+from reaf.simulation.torch._src.common.free_device import FreeDevice
+from reaf.simulation.torch._src.common.free_device import MocapDevice
+from reaf.simulation.torch._src.common.functional_device import joint
+from reaf.simulation.torch._src.common.functional_device import joint_pos
+from reaf.simulation.torch._src.common.functional_device import joint_vel
+from reaf.simulation.torch._src.common.functional_device import MeasurementDevice
+from reaf.simulation.torch._src.common.functional_observation import FunctionalObservationAdapter
+from reaf.simulation.torch._src.common.functional_reset import FunctionalEnvironmentReset
+from reaf.simulation.torch._src.common.functional_reset import TYPE_KEY
+from reaf.simulation.torch._src.common.functional_reset import TYPE_RESET
+from reaf.simulation.torch._src.common.functional_reset import TYPE_STARTUP
+from reaf.simulation.torch._src.common.functional_reward import FunctionalRewardProvider
+from reaf.simulation.torch._src.common.functional_reward import PassthroughRewardProvider
+from reaf.simulation.torch._src.common.functional_termination import FunctionalTerminationChecker
+from reaf.simulation.torch._src.common.functional_termination import TerminationResult
+from reaf.simulation.torch._src.common.inside_site_device import InsideSiteDevice
+from reaf.simulation.torch._src.common.model_randomization import randomize_actuator_gains
+from reaf.simulation.torch._src.common.model_randomization import randomize_field
+from reaf.simulation.torch._src.common.model_randomization import randomize_quat
+from reaf.simulation.torch._src.common.mujoco_camera_device import MujocoCameraDevice
+from reaf.simulation.torch._src.common.mujoco_camera_device import RenderMode
+from reaf.simulation.torch._src.common.multi_asset_device import MultiAsset
+from reaf.simulation.torch._src.common.multi_asset_device import MultiAssetDevice
+from reaf.simulation.torch._src.common.processors import DeltaActions
+from reaf.simulation.torch._src.common.processors import StochasticSubstepCommands
+from reaf.simulation.torch._src.common.processors import StochasticSubstepMeasurements
+from reaf.simulation.torch._src.common.processors import TransformActions
+from reaf.simulation.torch._src.common.resetter import CompositeResetter
+from reaf.simulation.torch._src.common.resetter import DomainRandomizationResetter
+from reaf.simulation.torch._src.common.resetter import NanTerminationResetter
+from reaf.simulation.torch._src.common.resetter import TerminationResetter
+from reaf.simulation.torch._src.common.site_device import SiteDevice
+from reaf.simulation.torch._src.common.spec_helper import SpecHelper
+from reaf.simulation.torch._src.common.spec_helper import SpecTemplate
+from reaf.simulation.torch._src.common.timeout_termination_checker import TimeoutTerminationChecker
+from reaf.simulation.torch._src.common.zero_reward_provider import ZeroRewardProvider
+from reaf.simulation.torch._src.core import clock
+from reaf.simulation.torch._src.core import device
+from reaf.simulation.torch._src.core import device_layer
+from reaf.simulation.torch._src.core import discount_provider
+from reaf.simulation.torch._src.core import environment
+from reaf.simulation.torch._src.core import feature_producer
+from reaf.simulation.torch._src.core import mujoco_types
+from reaf.simulation.torch._src.core import observation_space_adapter
+from reaf.simulation.torch._src.core import profile
+from reaf.simulation.torch._src.core import protocols
+from reaf.simulation.torch._src.core import reward_provider
+from reaf.simulation.torch._src.core import simulator
+from reaf.simulation.torch._src.core import substep_commands_processor
+from reaf.simulation.torch._src.core import substep_measurements_processor
+from reaf.simulation.torch._src.core import task_layer
+from reaf.simulation.torch._src.core import termination_checker
+from reaf.simulation.torch._src.core import torch_wrapper
+from reaf.simulation.torch._src.core import types
+from reaf.simulation.torch._src.core.clock import Clock
+from reaf.simulation.torch._src.core.device import Device
+from reaf.simulation.torch._src.core.device_layer import DeviceLayer
+from reaf.simulation.torch._src.core.discount_provider import DiscountProvider
+from reaf.simulation.torch._src.core.environment import AutoResetMode
+from reaf.simulation.torch._src.core.environment import EndOfEpisodeHandler
+from reaf.simulation.torch._src.core.environment import Environment
+from reaf.simulation.torch._src.core.environment import EnvironmentReset
+from reaf.simulation.torch._src.core.feature_producer import FeatureProducer
+from reaf.simulation.torch._src.core.mujoco_types import get_spec_by_obj_type
+from reaf.simulation.torch._src.core.mujoco_types import match_names
+from reaf.simulation.torch._src.core.mujoco_types import OBJ_TYPE_TO_ACCESSOR
+from reaf.simulation.torch._src.core.mujoco_types import PREFIX_TO_OBJ_TYPE
+from reaf.simulation.torch._src.core.mujoco_types import resolve_names
+from reaf.simulation.torch._src.core.mujoco_types import resolve_values
+from reaf.simulation.torch._src.core.mujoco_types import soft_joint_limits
+from reaf.simulation.torch._src.core.mujoco_types import SPEC_TYPE_INFO
+from reaf.simulation.torch._src.core.observation_space_adapter import ObservationSpaceAdapter
+from reaf.simulation.torch._src.core.protocols import TorchDeviceLayer
+from reaf.simulation.torch._src.core.reward_provider import RewardProvider
+from reaf.simulation.torch._src.core.reward_provider import RewardSpec
+from reaf.simulation.torch._src.core.reward_provider import RewardValue
+from reaf.simulation.torch._src.core.simulator import Simulator
+from reaf.simulation.torch._src.core.substep_commands_processor import SubstepCommandsProcessor
+from reaf.simulation.torch._src.core.substep_measurements_processor import SubstepMeasurementsProcessor
+from reaf.simulation.torch._src.core.task_layer import TaskLayer
+from reaf.simulation.torch._src.core.termination_checker import TerminationChecker
+from reaf.simulation.torch._src.core.torch_wrapper import JointType
+from reaf.simulation.torch._src.core.torch_wrapper import MjField
+from reaf.simulation.torch._src.core.torch_wrapper import MjObjType
+from reaf.simulation.torch._src.core.torch_wrapper import SetConstLevel
+from reaf.simulation.torch._src.core.torch_wrapper import TorchTensor
+from reaf.simulation.torch._src.core.torch_wrapper import TorchWrapper
+from reaf.simulation.torch._src.core.types import ArrayType
+from reaf.simulation.torch._src.core.types import Data
+from reaf.simulation.torch._src.core.types import Impl
+from reaf.simulation.torch._src.core.types import Model
+from reaf.simulation.torch._src.core.types import NUMPY_TO_TORCH_DTYPE
+from reaf.simulation.torch._src.core.types import Term
+from reaf.simulation.torch._src.core.types import TORCH_TO_NUMPY_DTYPE

reaf/simulation/torch/_src/builders/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Sim builders package."""
+
+from reaf.simulation.torch._src.builders import flat_scene as flat_scene

reaf/simulation/torch/_src/builders/config.py

@@ -0,0 +1,69 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared configuration for REAF environment builders."""
+
+from dataclasses import dataclass
+from typing import Literal
+
+import mujoco
+
+
+@dataclass
+class Sim:
+  """Simulation parameters."""
+
+  nworld: int = 1
+  xpu_device: str = "cuda:0"
+  impl: Literal["warp", "cpp"] = "warp"
+  control_timestep: float = 0.02
+  simulation_timestep: float = 0.005
+  jacobian: mujoco.mjtJacobian | None = None
+  integrator: mujoco.mjtIntegrator | None = None
+  iterations: int | None = None
+  ls_iterations: int | None = None
+  impratio: float | None = None
+  cone: mujoco.mjtCone | None = None
+  solver: mujoco.mjtSolver | None = None
+  gravity: tuple[float, float, float] | None = None
+  ccd_iterations: int | None = None
+  # warp only
+  nconmax: int | None = None
+  njmax: int | None = None
+  nccdmax: int | None = None
+  nvmax: int | None = None
+
+
+def apply_sim_config(spec: mujoco.MjSpec, sim: Sim) -> mujoco.MjSpec:
+  """Applies physics and solver parameters from Sim config to an MjSpec."""
+  spec.option.timestep = sim.simulation_timestep
+  if sim.jacobian is not None:
+    spec.option.jacobian = sim.jacobian
+  if sim.integrator is not None:
+    spec.option.integrator = sim.integrator
+  if sim.iterations is not None:
+    spec.option.iterations = sim.iterations
+  if sim.ls_iterations is not None:
+    spec.option.ls_iterations = sim.ls_iterations
+  if sim.impratio is not None:
+    spec.option.impratio = sim.impratio
+  if sim.cone is not None:
+    spec.option.cone = sim.cone
+  if sim.solver is not None:
+    spec.option.solver = sim.solver
+  if sim.gravity is not None:
+    spec.option.gravity = sim.gravity
+  if sim.ccd_iterations is not None:
+    spec.option.ccd_iterations = sim.ccd_iterations
+  return spec

reaf/simulation/torch/_src/builders/config_test.py

@@ -0,0 +1,55 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Sim config and helper utilities."""
+
+from absl.testing import absltest
+import mujoco
+from reaf.simulation.torch._src.builders import config
+
+
+class ConfigTest(absltest.TestCase):
+
+  def test_apply_sim_config(self):
+    mj_spec = mujoco.MjSpec()
+    sim = config.Sim(
+        iterations=12,
+        ls_iterations=18,
+        impratio=2.0,
+        solver=mujoco.mjtSolver.mjSOL_CG,
+        jacobian=mujoco.mjtJacobian.mjJAC_SPARSE,
+        cone=mujoco.mjtCone.mjCONE_ELLIPTIC,
+    )
+    result = config.apply_sim_config(mj_spec, sim)
+    self.assertIs(result, mj_spec)
+    self.assertEqual(mj_spec.option.iterations, 12)
+    self.assertEqual(mj_spec.option.ls_iterations, 18)
+    self.assertAlmostEqual(mj_spec.option.impratio, 2.0)
+    self.assertEqual(mj_spec.option.solver, mujoco.mjtSolver.mjSOL_CG)
+    self.assertEqual(mj_spec.option.jacobian, mujoco.mjtJacobian.mjJAC_SPARSE)
+    self.assertEqual(mj_spec.option.cone, mujoco.mjtCone.mjCONE_ELLIPTIC)
+
+  def test_apply_sim_config_preserves_unspecified_options(self):
+    mj_spec = mujoco.MjSpec()
+    mj_spec.option.iterations = 77
+    mj_spec.option.impratio = 3.3
+    sim = config.Sim()
+    result = config.apply_sim_config(mj_spec, sim)
+    self.assertIs(result, mj_spec)
+    self.assertEqual(mj_spec.option.iterations, 77)
+    self.assertAlmostEqual(mj_spec.option.impratio, 3.3)
+
+
+if __name__ == "__main__":
+  absltest.main()

reaf/simulation/torch/_src/builders/flat_scene/__init__.py

@@ -0,0 +1,31 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Flat scene builder package."""
+
+from reaf.simulation.torch._src.builders.config import apply_sim_config
+from reaf.simulation.torch._src.builders.config import Sim
+from reaf.simulation.torch._src.builders.flat_scene.build import build_device_layer
+from reaf.simulation.torch._src.builders.flat_scene.build import build_env
+from reaf.simulation.torch._src.builders.flat_scene.build import build_scene
+from reaf.simulation.torch._src.builders.flat_scene.build import build_task_layer
+from reaf.simulation.torch._src.builders.flat_scene.build import make_deployment_config
+from reaf.simulation.torch._src.builders.flat_scene.config import Component
+from reaf.simulation.torch._src.builders.flat_scene.config import DeviceLayer
+from reaf.simulation.torch._src.builders.flat_scene.config import Environment
+from reaf.simulation.torch._src.builders.flat_scene.config import Obs
+from reaf.simulation.torch._src.builders.flat_scene.config import ObsGroup
+from reaf.simulation.torch._src.builders.flat_scene.config import Reward
+from reaf.simulation.torch._src.builders.flat_scene.config import TaskLayer
+from reaf.simulation.torch._src.builders.flat_scene.config import Term