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

[jeff-hykin]

• . Only --simulation changes.

It can't be a bool cause 1 robot can have many sims

[mustafab0]

Ah. I wanted to keep this as a simple toggle between sim or hw. And most of the configs just go into a SimConfig to define the engine etc.

Otherwise we have to maintain simulation argument with many other kwargs to track. I dont feel very strongly about this.

[hvent90]

What is the shape of WorldState and SimTruth? My initial interpretation of this is that WorldState would already have everything needed for evals.

Same question applies for RobotState, RobotCommand, and what some example usages of SimulationModule.apply would be.

Common useful operations in sims:

• querying actors and getting their transforms and physics state (velocity, angular torque)
• allowing for callbacks to be registered for physics events (ex: on collision)
• spawning actors (such as chairs/tables/cups)

Something that evals do is take a pre-made scene, such as an apartment, and spawn an object such as a chair, and then the chair will be used in the eval case.

Common APIs that I've ran into that satisfies most needs for consumers:

class Actor:
    transform: (vector3, quaternion)
    # I forget how physics is stored but probably a property for velocity and angular torque
    def set_world_transform(transform):
    def get_world_transform(): 
    def on_collision(OnCollisionHandler):
        OnCollisionHandler(physics_info) # the thing it collided with, point of impact, force of impact, etc

class SimInstance: # this is equivalent to SimulationModule and would handle the actual life cycle of the child process 
    def load_world(World): # load scene

class World: # this is the scene and responsible for all objects as well as ticking physics
    def spawn_actor(ActorClass, transform): # this would either take a callback to allow the caller to tweak specific params of the actor or it can take a config value that would achieve the same
    def delete_actor(ActorReference):

[mustafab0]

What is the shape of WorldState and SimTruth? My initial interpretation of this is that WorldState would already have everything needed for evals.

1. world_state is a Sim's raw state, a list of numbers. Sensor workers copy it into their own model to render. You can't read it without the model file.
2. sim_truth is readable and organised by name: each object's pose, velocity and size, the contacts, and the robot poses. This is what eval should use. we can ofcourse rename it.

Same question applies for RobotState, RobotCommand, and what some example usages of SimulationModule.apply would be.

This would be some examples of how you would use it.

RobotCommand = MotorTargets | JointPositions | Velocity | Behavior   # one per SimRobotClient call
RobotState   # joints (position, velocity, effort), IMU, base pose and speed

sim.apply("g1",    MotorTargets(q=..., dq=..., kp=..., kd=..., tau=...))  # MuJoCo: motor torque per joint
sim.apply("arm",   JointPositions({"arm/joint1": 0.3}))                   # MuJoCo: joint position targets
sim.apply("go2_a", Velocity(Twist(linear=(0.5, 0, 0))))                 # target speed for the walking policy

[Nabla7]

Base of the current draft sim2 stack: #4723