Skip to main content

​
rfx.sim

Simulation backends for parallel training with support for Genesis (GPU), MuJoCo MJX (JAX), and Mock (CPU testing).

​
SimRobot

Simulated robot with parallel environments. 
from rfx.sim import SimRobot

robot = SimRobot.from_config("so101.yaml", num_envs=4096, backend="genesis")
obs = robot.observe()       # (4096, 64)
robot.act(action)           # Step all 4096 envs

​
Constructor

SimRobot(
    config: str | Path | RobotConfig | dict,
    num_envs: int = 1,
    backend: str = "mock",
    device: str = "cuda",
    **kwargs
)
​
config
str | Path | RobotConfig | dict
required
Robot configuration. Can be a path to a YAML file, a RobotConfig object, or a dictionary.
​
num_envs
int
default:"1"
Number of parallel environments to simulate.
​
backend
str
default:"mock"
Simulation backend to use. Options: "genesis", "mjx", "mock".
​
device
str
default:"cuda"
Device to run simulation on ("cuda" or "cpu").
​
**kwargs
dict
Additional backend-specific arguments.

​
Methods

​
from_config

@classmethod
SimRobot.from_config(
    config_path: str | Path,
    num_envs: int = 1,
    backend: str = "mock",
    device: str | None = None,
    **kwargs
) -> SimRobot
Create a SimRobot from a configuration file.
​
config_path
str | Path
required
Path to robot configuration YAML file.
​
num_envs
int
default:"1"
Number of parallel environments.
​
backend
str
default:"mock"
Simulation backend to use.
​
device
str | None
default:"None"
Device to use. If None, automatically selects CUDA if available, otherwise CPU.
​
return
SimRobot
Initialized SimRobot instance.

​
observe

robot.observe() -> dict[str, torch.Tensor]
Get current observations from all environments.
​
return
dict[str, torch.Tensor]
Dictionary containing observation tensors with keys like "state", "images", etc.

​
act

robot.act(action: torch.Tensor) -> None
Apply actions to all environments and step the simulation.
​
action
torch.Tensor
required
Action tensor of shape (num_envs, action_dim).

​
reset

robot.reset(env_ids: torch.Tensor | None = None) -> dict[str, torch.Tensor]
Reset specified environments (or all if env_ids is None).
​
env_ids
torch.Tensor | None
default:"None"
Indices of environments to reset. If None, resets all environments.
​
return
dict[str, torch.Tensor]
Initial observations after reset.

​
get_reward

robot.get_reward() -> torch.Tensor
Get rewards for all environments.
​
return
torch.Tensor
Reward tensor of shape (num_envs,).

​
get_done

robot.get_done() -> torch.Tensor
Get done flags for all environments.
​
return
torch.Tensor
Boolean tensor of shape (num_envs,) indicating which environments are done.

​
render

robot.render() -> None
Render the simulation (backend-dependent).

​
close

robot.close() -> None
Clean up simulation resources.

​
Properties

​
backend
str
Name of the simulation backend being used.
​
config
RobotConfig
Robot configuration object.

​
MockRobot

Mock robot for testing without real simulation dependencies. 
from rfx.sim import MockRobot

robot = MockRobot(state_dim=12, action_dim=6, num_envs=8)
obs = robot.observe()
robot.act(torch.zeros(8, 6))

​
Constructor

MockRobot(
    state_dim: int = 12,
    action_dim: int = 6,
    num_envs: int = 1,
    max_state_dim: int = 64,
    max_action_dim: int = 64,
    device: str = "cpu",
    **kwargs
)
​
state_dim
int
default:"12"
Dimension of the state space.
​
action_dim
int
default:"6"
Dimension of the action space.
​
num_envs
int
default:"1"
Number of parallel environments.
​
max_state_dim
int
default:"64"
Maximum state dimension for padding.
​
max_action_dim
int
default:"64"
Maximum action dimension for padding.
​
device
str
default:"cpu"
Device to run on.
​
**kwargs
dict
Additional configuration options (e.g., max_steps).

​
Methods

MockRobot implements the same interface as SimRobot:
  • observe() -> dict[str, torch.Tensor]
  • act(action: torch.Tensor) -> None
  • reset(env_ids: torch.Tensor | None = None) -> dict[str, torch.Tensor]
  • get_reward() -> torch.Tensor
  • get_done() -> torch.Tensor

Build docs developers (and LLMs) love

Get started for freeTalk to us