Refactor: sim2, a standard Sim interface for DimOS

[mustafab0]

Refactor: sim2, a standard Sim interface for DimOS

Builds on Pim's sim2 (pim/feat/sim2-core)

Goals

• The same blueprint runs on the real robot and in sim. Only --simulation changes.
• Namespaced robots share a sim and can see each other.
• The sim knows no vendor protocol. Connections talk to it through one client.

Overview

flowchart TB
  
  CC["Nav, , agents, ControlCoordinator<br/>arbitrates"] --> CONN["Connection module<br/>one per robot interface"]
  CONN -- "real" --> SDK["Vendor SDK"] --> ROBOT["Robot"]
  CONN -- "sim" --> SA["Sim adapter"] --> CLIENT["SimRobotClient"]
  CLIENT -- "Zenoh, robot id" --> SIM["MujocoSimModule, DimSimModule, ...<br/>one per engine"]
  SIM -- "world_state" --> SEN["Sim sensor modules"]
  SEN -- "same topics as real sensors" --> CC
  SIM -- "sim_truth" --> EVAL["Evals"]

Loading

Connection module: Every connection module also provides a SimAdapter

class G1LowLevelConnection(ConnectionModule):                   
    motor_command: In[MotorCommandArray]
    joint_state: Out[JointState]
    imu: Out[Imu]

    def describe(self) -> ControlDescription:                   
        return pd_joint_description("g1", G1_JOINTS, model=G1_MODEL)

    def describe_simulated_robot(self) -> SimRobot:                # for the sim
        return SimRobot(model=G1_MJCF, mount=self.config.mount)

    def connect(self) -> None:
        if self.config.simulation:
            self.hw = G1SimAdapter(SimRobotClient(self.robot_id))
        else:
            self.hw = UnitreeSdk(self.config.interface)

    def read_state(self):     return self.hw.read_state()
    def write(self, values):  self.hw.write(values)
    def halt(self):           self.hw.halt()


class G1SimAdapter:                                                # next to the driver
    def __init__(self, sim: SimRobotClient): self.sim = sim
    def read_state(self):     return self.sim.state()
    def write(self, values):  self.sim.set_motor_targets(values)   # q, dq, kp, kd, tau
    def halt(self):           self.sim.hold_position()

Go2 already works this way: make_connection() picks WebRTC, replay, MuJoCo or DimSim (go2/connection.py:143). Its MuJoCo and DimSim backends become one Go2SimConnection on SimRobotClient.

SimRobotClient: This is what translates between the SimModule and the robot ConnectionModule

class SimRobotClient:                       # one per robot; one Zenoh session; not a module
    def __init__(self, robot_id: str): ...

    def set_motor_targets(self, targets: dict[str, MotorTarget]) -> None: ...   # PD bodies
    def set_joint_positions(self, positions: dict[str, float]) -> None: ...     # arms
    def set_velocity(self, twist: Twist) -> None: ...                           # bases
    def hold_position(self) -> None: ...
    def state(self) -> RobotState: ...                                          # joints, imu, base pose

Startup: robots are collected from the blueprint

The engine can't add robots after it starts, so every robot is known before the sim starts.

sequenceDiagram
  participant B as Blueprint build
  participant C as Connection modules
  participant S as SimulationModule
  B->>C: describe_simulated_robot()
  C-->>B: SimRobot (model, mount)
  B->>S: start(scene, robots)
  S->>S: load scene + all robots, start physics
  C->>S: SimRobotClient connects by robot id

Loading

# the same blueprint for real and sim
g1 = autoconnect(
    G1LowLevelConnection.blueprint(),
    G1Coordinator.blueprint(tasks=[groot_task()]),
    HeadCamera.blueprint(),
)

# dimos run g1                                        real robot
# dimos --simulation mujoco --scene-package kitchen run g1    MujocoSimModule added automatically

SimulationModule: many robots, one world

flowchart LR
  subgraph BP["Blueprint: fleet"]
    A["Go2Connection<br/>go2_a"]
    B["Go2Connection<br/>go2_b"]
    X["XArmConnection<br/>arm"]
  end
  subgraph SIM["MujocoSimModule: one world"]
    W["kitchen scene<br/>go2_a, go2_b, arm"]
  end
  A -- "SimRobotClient go2_a" --> W
  B -- "SimRobotClient go2_b" --> W
  X -- "SimRobotClient arm" --> W
  W -- "world_state" --> SEN["Sensor workers<br/>cameras of all three robots"]

Loading

fleet = autoconnect(
    go2.namespace("go2_a"),
    go2.namespace("go2_b"),
    xarm.namespace("arm"),
)
# dimos --simulation mujoco --scene-package kitchen run fleet

# collected from the blueprint before the sim starts; robot id = namespace
robots = [
    SimRobot(id="go2_a", model=GO2_MJCF,   mount=Pose(0, 0, 0)),
    SimRobot(id="go2_b", model=GO2_MJCF,   mount=Pose(1, 0, 0)),
    SimRobot(id="arm",   model=XARM7_MJCF, mount=Pose(2, 0, 0.8)),
]
sim.load(scene="kitchen", robots=robots)                  # one world, all three robots

# each connection's sim adapter talks only to its own robot
SimRobotClient("go2_a").set_velocity(Twist(linear=(0.5, 0, 0)))
SimRobotClient("arm").set_joint_positions({"arm/joint1": 0.3})

One subclass per engine :

class SimulationModule(Module, ABC):            # the contract every engine keeps
    world_state: Out[WorldState]                # full state incl. joints
    sim_truth: Out[SceneState]                  # readable, for evals

    @rpc
    def reset(self, initial: SceneUpdate | None = None) -> SceneState: ...
    @rpc
    def set_scene_state(self, update: SceneUpdate) -> SceneState: ...
    @rpc
    def set_spawn(self, robot_id: str, xyz: Vec3, rpy: Vec3 = (0, 0, 0)) -> None: ...
    @rpc
    def set_paused(self, paused: bool) -> None: ...

    # each engine fills these in
    @abstractmethod
    def load(self, scene: Path, robots: list[SimRobot]) -> None: ...
    @abstractmethod
    def apply(self, robot_id: str, command: RobotCommand) -> None: ...
    @abstractmethod
    def step(self) -> None: ...
    @abstractmethod
    def robot_state(self, robot_id: str) -> RobotState: ...


class MujocoSimModule(SimulationModule): ...    # starts camera and lidar workers
class DimSimModule(SimulationModule): ...
class IsaacSimModule(SimulationModule): ...     # renders cameras inside Isaac

dimos --simulation mujoco run g1     # adds MujocoSimModule
dimos --simulation dimsim run g1     # adds DimSimModule

Only real engines get a subclass. RoboCasa and LIBERO run on MuJoCo, so they are scenes for MujocoSimModule.

Transport

• world_state, sim_truth and robot state are normal Zenoh traffic.
• Delete sim2/ipc/. Default to Zenoh as it supports Shm naturally.

Rules

• Connection RPCs are driver primitives (change mode, set a register).
• A sim adapter fakes vendor behaviour the vendor doesn't give us (e.g. sit_down).
• The sim starts each robot the way the OEM boots it.

[paul-nechifor]

What does "same topics as real sensors" mean? Would MujocoSimModule publish directly to /go2_a/color_image, /go2_b/color_image, etc. (If so, there are some problems with doing it that way.)

Love the sequence diagram for startup. I think one important one would be do describe how communication from the connection module ends up on the simulator.

For example, let's say the simulator has a liedown policy implemented. Let's say someone calls the rpc method GO2Connection.liedown. That calls MujocoConnection.liedown. Currently that does nothing. But how would you talk to the Mujoco sim to do it for that robot?

[mustafab0]

What does "same topics as real sensors" mean? Would MujocoSimModule publish directly to /go2_a/color_image, /go2_b/color_image, etc. (If so, there are some problems with doing it that way.)

No, the Go2ConnectionModule in this case would publish camera images. It will get the camera renders through the SimRobotClient. So the images will be namespaced as per the ConnectionModule

The following is an end to end command example:

sequenceDiagram
  participant U as Caller
  participant GC as GO2Connection go2_a
  participant GS as Go2SimConnection
  participant SC as SimRobotClient go2_a
  participant SM as MujocoSimModule
  U->>GC: liedown() RPC
  GC->>GS: liedown()
  GS->>SC: behavior(liedown)
  SC->>SM: Zenoh query sim/go2_a/behavior/liedown
  SM->>SM: start go2_a LiedownPolicy
  SM-->>SC: reply ok
  SC-->>GC: True
  loop every physics step until done
    SM->>SM: LiedownPolicy drives go2_a, step
  end

Loading