xwm.envs¶

A Franka Emika FR3 arm in Newton, observed as pixels or as a 20-D proprioceptive state, with a dense reach reward. Needs the newton extra.

Simulated environments with real dynamics.

Distinct from xwm.data, which generates synthetic arrays in pure JAX. Everything here wraps an external simulator, carries optional heavy dependencies, and lives outside the jit boundary.

Modules:

Name	Description
`discretize`	Turning a continuous action space into a discrete one.
`newton_franka`	A Franka arm in Newton, as a source of action-labelled video.
`render`	Two rendering paths, for two different jobs.

Classes:

Name	Description
`FrankaConfig`	Simulation and rendering settings.
`FrankaEnv`	A Franka FR3 arm with RGB observations and joint-space actions.
`HighQualityRenderer`	Render or export a trajectory with Newton's high-quality viewers.

Functions:

Name	Description
`discrete_action_table`	`(n_actions, action_dim)` table of axis-aligned moves.
`franka_sequences`	Collect an action-labelled video dataset from `env`.
`smooth_actions`	`(n, length, action_dim)` temporally correlated actions in `[-1, 1]`.
`look_at_angles`	`(pitch, yaw)` in degrees for a viewer camera at `eye` facing `target`.
`supersample`	Box-downsample a `(3, Hf, Wf)` frame by `factor`.
`which_backends`	Which rendering backends are usable here.

FrankaConfig ¶

FrankaConfig(image_size: int = 64, action_scale: float = 0.25, fps: int = 30, substeps: int = 16, joint_armature: float = 0.1, target_ke: float = 800.0, target_kd: float = 40.0, pose_noise: float = 0.35, goal: tuple[float, float, float] = (0.3, 0.15, 0.55), action_penalty: float = 0.01, camera_distance: float = 1.15, camera_height: float = 0.65, camera_target_height: float = 0.45, camera_fov_degrees: float = 60.0, enable_shadows: bool = False, enable_textures: bool = False, solver: str = 'auto', asset_name: str = 'franka_emika_panda', urdf_relative_path: str = 'urdf/fr3_franka_hand.urdf')

Simulation and rendering settings.

Parameters:

Name	Type	Description	Default
`image_size`	`int`	square observation resolution. CPU raytracing cost grows with the pixel count, so 64 is a sensible default for training.	`64`
`action_scale`	`float`	radians of joint-target change per unit action. Large enough that one step is visible in the image, which matters: if a single step barely changes the observation, "predict no change" becomes a near-perfect baseline and the benchmark is vacuous.	`0.25`
`fps`	`int`	control rate.	`30`
`substeps`	`int`	physics substeps per control step. What matters is the resulting `dt / substeps`; Featherstone integration of this arm diverges above roughly 1/480 s.	`16`
`joint_armature`	`float`	added rotor inertia. The single most important stabiliser here -- without it the solver produces NaNs.	`0.1`
`target_ke, target_kd`		joint position-servo gains. The defaults were chosen so the arm holds its commanded pose (~0.007 rad of drift over 12 idle steps) while a commanded sweep still moves the tool ~0.65 m. Both halves matter: a servo too weak and the model learns gravity instead of the action; too stiff and the solver diverges.	required
`pose_noise`	`float`	radians of uniform noise on the initial pose at reset, which is what gives the dataset its variety.	`0.35`
`enable_shadows, enable_textures`		raytracing quality. Both cost render time, which is why they are off by default on CPU; on a GPU they are close to free and they add real information to the image -- shading disambiguates depth, and texture distinguishes links that are otherwise identically white.	required
`solver`	`str`	`"mujoco"`, `"featherstone"`, or `"auto"` to prefer MuJoCo (GPU-only, via `mujoco_warp`) and fall back to Featherstone.	`'auto'`

FrankaEnv ¶

FrankaEnv(config: FrankaConfig | None = None, *, urdf_path: str | Path | None = None)

A Franka FR3 arm with RGB observations and joint-space actions.

Example

>>> env = FrankaEnv()
>>> env.reset(seed=0)
>>> frame = env.observe()               # (3, H, W) in [0, 1]
>>> env.step(np.zeros(env.action_dim))  # joint-target deltas in [-1, 1]

Attributes:

Name	Type	Description
`config`		the `FrankaConfig` in use.
`model`		the underlying `newton.Model`, for callers who need it.

Methods:

Name	Description
`reset`	Reset to a randomly perturbed home pose. Returns the observation.
`step`	Apply one control step. Returns the resulting observation.
`observe`	Render the current state to `(3, H, W)` float32 in `[0, 1]`.
`render`	Render at any resolution, `(3, size, size)` float32 in `[0, 1]`.
`joint_positions`	`(7,)` arm joint angles. For evaluation -- the model never sees these.
`body_positions`	`(n_bodies, 3)` world-frame body origins.
`tool_position`	`(3,)` world position of the tool centre point (the hand).
`state_observation`	`(20,)` proprioceptive observation: joints, velocities, tool, goal delta.
`is_finite`	Whether the simulation is still numerically healthy.
`high_quality_renderer`	A `xwm.envs.HighQualityRenderer` bound to this environment's model.
`goal_distance`	Metres from the tool to the goal. Ground truth, for evaluation.
`reward`	Dense reach reward in roughly `[-1, 0]`, minus an action penalty.
`rollout`	Execute an action sequence from a fresh reset.

Source code in xwm/envs/newton_franka.py

def __init__(self, config: FrankaConfig | None = None, *, urdf_path: str | Path | None = None):
    newton, wp = _require_newton()
    self.config = config or FrankaConfig()
    self._newton, self._wp = newton, wp

    if urdf_path is None:
        import newton.utils

        asset = newton.utils.download_asset(self.config.asset_name)
        urdf_path = Path(asset) / self.config.urdf_relative_path
    urdf_path = Path(urdf_path)
    if not urdf_path.exists():
        raise FileNotFoundError(f"Franka URDF not found at {urdf_path}")

    builder = newton.ModelBuilder()
    builder.default_joint_cfg = newton.ModelBuilder.JointDofConfig(
        armature=self.config.joint_armature,
        limit_ke=1.0e4,
        limit_kd=1.0e2,
    )
    builder.add_urdf(str(urdf_path), floating=False)
    builder.add_ground_plane()
    # Set the servo gains on the builder arrays directly. Passing them
    # through JointDofConfig does *not* reach the arm joints in Newton 1.5 --
    # target_kd silently stays 0, which makes the position servo an
    # undamped spring and sends the Featherstone solver to NaN.
    for i in range(len(builder.joint_target_ke)):
        builder.joint_target_ke[i] = self.config.target_ke
        builder.joint_target_kd[i] = self.config.target_kd
    self.model = builder.finalize()
    self._verify_gains()

    self.solver, self.solver_name = self._make_solver()
    self._state_0 = self.model.state()
    self._state_1 = self.model.state()
    self._control = self.model.control()

    self._n_dofs = int(self.model.joint_dof_count)
    limits_lower = np.asarray(self.model.joint_limit_lower.numpy(), dtype=np.float32)
    limits_upper = np.asarray(self.model.joint_limit_upper.numpy(), dtype=np.float32)
    self.joint_limit_lower = limits_lower[: self._n_dofs]
    self.joint_limit_upper = limits_upper[: self._n_dofs]

    self.tool_body_index = self._find_tool_body()
    self._setup_camera()
    self._target = np.zeros(self._n_dofs, dtype=np.float32)
    self.reset(seed=0)

camera_framing ¶

camera_framing: tuple[ndarray, ndarray]

(eye, target) in world metres, shared by every renderer.

One definition so that a path-traced figure and the observations the model trains on show the same view from the same place.

action_dim ¶

action_dim: int

Number of controlled arm joints (the gripper is held fixed).

state_dim ¶

state_dim: int

Length of state_observation: joints, velocities, tool, delta.

state ¶

state

The current newton.State, for a renderer or a custom sensor.

goal ¶

goal: ndarray

(3,) reach target for the tool.

reset ¶

reset(seed: int | None = None) -> ndarray

Reset to a randomly perturbed home pose. Returns the observation.

Source code in xwm/envs/newton_franka.py

def reset(self, seed: int | None = None) -> np.ndarray:
    """Reset to a randomly perturbed home pose. Returns the observation."""
    newton, wp = self._newton, self._wp
    rng = np.random.default_rng(seed)
    q = np.zeros(self._n_dofs, dtype=np.float32)
    n = min(N_ARM_JOINTS, self._n_dofs)
    q[:n] = HOME_POSE[:n] + rng.uniform(
        -self.config.pose_noise, self.config.pose_noise, size=n
    ).astype(np.float32)
    q = np.clip(q, self.joint_limit_lower, self.joint_limit_upper)

    self._state_0.joint_q.assign(wp.array(q, dtype=wp.float32))
    self._state_0.joint_qd.assign(wp.zeros(self._n_dofs, dtype=wp.float32))
    newton.eval_fk(self.model, self._state_0.joint_q, self._state_0.joint_qd, self._state_0)
    self._target = q.copy()
    self._control.joint_target_q = wp.array(self._target, dtype=wp.float32)
    return self.observe()

step ¶

step(action: ndarray) -> ndarray

Apply one control step. Returns the resulting observation.

Parameters:

Name	Type	Description	Default
`action`	`ndarray`	`(7,)` in `[-1, 1]`; interpreted as a delta on the joint position targets, scaled by `config.action_scale` and clipped to the URDF joint limits. Deltas rather than absolute targets so the action distribution is state-independent, which is what makes a learned `(z, a) -> z'` well-posed.	required

Source code in xwm/envs/newton_franka.py

def step(self, action: np.ndarray) -> np.ndarray:
    """Apply one control step. Returns the resulting observation.

    Args:
        action: ``(7,)`` in ``[-1, 1]``; interpreted as a *delta* on the
            joint position targets, scaled by ``config.action_scale`` and
            clipped to the URDF joint limits. Deltas rather than absolute
            targets so the action distribution is state-independent, which
            is what makes a learned ``(z, a) -> z'`` well-posed.
    """
    wp = self._wp
    action = np.clip(np.asarray(action, dtype=np.float32).reshape(-1), -1.0, 1.0)
    if action.shape[0] != self.action_dim:
        raise ValueError(f"expected {self.action_dim} action dims, got {action.shape[0]}")

    n = min(N_ARM_JOINTS, self._n_dofs)
    self._target[:n] = np.clip(
        self._target[:n] + self.config.action_scale * action[:n],
        self.joint_limit_lower[:n],
        self.joint_limit_upper[:n],
    )
    self._control.joint_target_q = wp.array(self._target, dtype=wp.float32)

    sub_dt = self.dt / self.config.substeps
    for _ in range(self.config.substeps):
        self._state_0.clear_forces()
        self.solver.step(self._state_0, self._state_1, self._control, None, sub_dt)
        self._state_0, self._state_1 = self._state_1, self._state_0
    return self.observe()

observe ¶

observe() -> ndarray

Render the current state to (3, H, W) float32 in [0, 1].

Uses config.image_size -- the resolution the model is trained on.

Source code in xwm/envs/newton_franka.py

def observe(self) -> np.ndarray:
    """Render the current state to ``(3, H, W)`` float32 in ``[0, 1]``.

    Uses ``config.image_size`` -- the resolution the model is trained on.
    """
    return self.render()

render ¶

render(size: int | None = None, *, samples: int = 1) -> ndarray

Render at any resolution, (3, size, size) float32 in [0, 1].

Pass a larger size for figures: upscaling a training frame turns every pixel into a block and adds no detail, while the raytracer will render at whatever resolution you ask for.

Parameters:

Name	Type	Description	Default
`size`	`int \| None`	output resolution. Defaults to `config.image_size`.	`None`
`samples`	`int`	supersampling factor. The Warp raytracer casts one ray per pixel, so silhouettes and shadow boundaries come out as hard staircases; rendering at `samples x` and averaging down fixes that at `samples ** 2` the cost. Leave at 1 for training observations, use 2-3 for figures.	`1`

For photorealistic output -- soft shadows, ambient occlusion, materials -- see xwm.envs.HighQualityRenderer, which drives Newton's OVRTX path tracer or exports a USD stage.

Source code in xwm/envs/newton_franka.py

def render(self, size: int | None = None, *, samples: int = 1) -> np.ndarray:
    """Render at any resolution, ``(3, size, size)`` float32 in ``[0, 1]``.

    Pass a larger ``size`` for figures: upscaling a training frame turns
    every pixel into a block and adds no detail, while the raytracer will
    render at whatever resolution you ask for.

    Args:
        size: output resolution. Defaults to ``config.image_size``.
        samples: supersampling factor. The Warp raytracer casts one ray per
            pixel, so silhouettes and shadow boundaries come out as hard
            staircases; rendering at ``samples x`` and averaging down fixes
            that at ``samples ** 2`` the cost. Leave at 1 for training
            observations, use 2-3 for figures.

    For photorealistic output -- soft shadows, ambient occlusion, materials --
    see :class:`xwm.envs.HighQualityRenderer`, which drives Newton's OVRTX
    path tracer or exports a USD stage.
    """
    size = self.config.image_size if size is None else int(size)
    if samples < 1:
        raise ValueError(f"samples must be at least 1, got {samples}")
    if samples > 1:
        from .render import supersample

        return supersample(self._render_once(size * samples), samples)
    return self._render_once(size)

joint_positions ¶

joint_positions() -> ndarray

(7,) arm joint angles. For evaluation -- the model never sees these.

Source code in xwm/envs/newton_franka.py

def joint_positions(self) -> np.ndarray:
    """``(7,)`` arm joint angles. For evaluation -- the model never sees these."""
    q = np.array(self._state_0.joint_q.numpy(), dtype=np.float32, copy=True)
    return q[: self.action_dim]

body_positions ¶

body_positions() -> ndarray

(n_bodies, 3) world-frame body origins.

Source code in xwm/envs/newton_franka.py

def body_positions(self) -> np.ndarray:
    """``(n_bodies, 3)`` world-frame body origins."""
    q = np.array(self._state_0.body_q.numpy(), dtype=np.float32, copy=True)
    return q[:, :3]

tool_position ¶

tool_position() -> ndarray

(3,) world position of the tool centre point (the hand).

Source code in xwm/envs/newton_franka.py

def tool_position(self) -> np.ndarray:
    """``(3,)`` world position of the tool centre point (the hand)."""
    return np.array(self.body_positions()[self.tool_body_index], copy=True)

state_observation ¶

state_observation() -> ndarray

(20,) proprioceptive observation: joints, velocities, tool, goal delta.

The cheap alternative to pixels. TD-MPC2 and MuZero on state converge in minutes rather than hours, which makes them testable; swap in an image encoder once the pipeline is known to work.

Source code in xwm/envs/newton_franka.py

def state_observation(self) -> np.ndarray:
    """``(20,)`` proprioceptive observation: joints, velocities, tool, goal delta.

    The cheap alternative to pixels. TD-MPC2 and MuZero on state converge in
    minutes rather than hours, which makes them testable; swap in an image
    encoder once the pipeline is known to work.
    """
    q = self.joint_positions()
    qd = np.array(self._state_0.joint_qd.numpy(), dtype=np.float32, copy=True)
    return np.concatenate(
        [q, qd[: self.action_dim], self.tool_position(), self.tool_position() - self.goal]
    ).astype(np.float32)

is_finite ¶

is_finite() -> bool

Whether the simulation is still numerically healthy.

Source code in xwm/envs/newton_franka.py

def is_finite(self) -> bool:
    """Whether the simulation is still numerically healthy."""
    return bool(np.all(np.isfinite(np.array(self._state_0.joint_q.numpy(), copy=True))))

high_quality_renderer ¶

high_quality_renderer(**kwargs)

A xwm.envs.HighQualityRenderer bound to this environment's model.

Feed it simulation states as the episode runs::

with env.high_quality_renderer(backend="usd",
                               output_path="episode.usd") as renderer:
    env.reset(seed=0)
    for action in actions:
        env.step(action)
        renderer.add(env.state)

Source code in xwm/envs/newton_franka.py

def high_quality_renderer(self, **kwargs):
    """A :class:`xwm.envs.HighQualityRenderer` bound to this environment's model.

    Feed it simulation states as the episode runs::

        with env.high_quality_renderer(backend="usd",
                                       output_path="episode.usd") as renderer:
            env.reset(seed=0)
            for action in actions:
                env.step(action)
                renderer.add(env.state)
    """
    from .render import HighQualityRenderer

    kwargs.setdefault("fps", self.config.fps)
    kwargs.setdefault("up_axis", "XYZ"[int(np.argmax(WORLD_UP))])
    renderer = HighQualityRenderer(self.model, **kwargs)
    renderer.look_at(*self.camera_framing)
    return renderer

goal_distance ¶

goal_distance() -> float

Metres from the tool to the goal. Ground truth, for evaluation.

Source code in xwm/envs/newton_franka.py

def goal_distance(self) -> float:
    """Metres from the tool to the goal. Ground truth, for evaluation."""
    return float(np.linalg.norm(self.tool_position() - self.goal))

reward ¶

reward(action: ndarray | None = None) -> float

Dense reach reward in roughly [-1, 0], minus an action penalty.

-tanh(distance) rather than -distance: a bounded reward keeps the value function inside the categorical head's bin range without per-task tuning, and its gradient does not vanish far from the goal the way a squared distance's does.

Source code in xwm/envs/newton_franka.py

def reward(self, action: np.ndarray | None = None) -> float:
    """Dense reach reward in roughly ``[-1, 0]``, minus an action penalty.

    ``-tanh(distance)`` rather than ``-distance``: a bounded reward keeps the
    value function inside the categorical head's bin range without per-task
    tuning, and its gradient does not vanish far from the goal the way a
    squared distance's does.
    """
    shaped = -float(np.tanh(self.goal_distance()))
    if action is None or self.config.action_penalty == 0.0:
        return shaped
    cost = self.config.action_penalty * float(np.mean(np.square(action)))
    return shaped - cost

rollout ¶

rollout(actions: ndarray, *, seed: int | None = None) -> dict[str, ndarray]

Execute an action sequence from a fresh reset.

Parameters:

Name	Type	Description	Default
`actions`	`ndarray`	`(T, 7)`.	required

Returns:

Type	Description
`dict[str, ndarray]`	``{"video": (T + 1, 3, H, W), "joint_q": (T + 1, 7),
`dict[str, ndarray]`	"tool": (T + 1, 3)}`` -- one more observation than actions, since
`dict[str, ndarray]`	the initial frame precedes the first action.

Source code in xwm/envs/newton_franka.py

def rollout(self, actions: np.ndarray, *, seed: int | None = None) -> dict[str, np.ndarray]:
    """Execute an action sequence from a fresh reset.

    Args:
        actions: ``(T, 7)``.

    Returns:
        ``{"video": (T + 1, 3, H, W), "joint_q": (T + 1, 7),
        "tool": (T + 1, 3)}`` -- one more observation than actions, since
        the initial frame precedes the first action.
    """
    actions = np.asarray(actions, dtype=np.float32)
    frames = [self.reset(seed=seed)]
    joints = [self.joint_positions()]
    tools = [self.tool_position()]
    states = [self.state_observation()]
    rewards = []
    for action in actions:
        frames.append(self.step(action))
        joints.append(self.joint_positions())
        tools.append(self.tool_position())
        states.append(self.state_observation())
        rewards.append(self.reward(action))
    return {
        "video": np.stack(frames),
        "joint_q": np.stack(joints),
        "tool": np.stack(tools),
        "state": np.stack(states),
        "reward": np.asarray(rewards, dtype=np.float32),
    }

HighQualityRenderer ¶

HighQualityRenderer(model, *, backend: Backend = 'usd', size: tuple[int, int] = (1280, 720), output_path: str | Path | None = None, environment: str = 'studio', fps: int = 30, up_axis: str = 'Z')

Render or export a trajectory with Newton's high-quality viewers.

Parameters:

Name	Type	Description	Default
`model`		the `newton.Model` to render.	required
`backend`	`Backend`	`"rtx"` for path-traced frames, `"usd"` for a USD stage.	`'usd'`
`size`	`tuple[int, int]`	output resolution (`"rtx"` only).	`(1280, 720)`
`output_path`	`str \| Path \| None`	destination stage (`"usd"` only).	`None`
`environment`	`str`	OVRTX lighting environment -- `"studio"` is the lit backdrop used for presentation renders.	`'studio'`
`fps`	`int`	playback rate recorded in the output.	`30`

Example

>>> renderer = HighQualityRenderer(env.model, backend="usd",
...                                output_path="episode.usd")
>>> for state in states:
...     renderer.add(state)
>>> renderer.close()

Methods:

Name	Description
`set_camera`	Forwarded to the viewer, where supported.
`look_at`	Aim the camera from `eye` at `target`, both in world metres.
`add`	Record one simulation state as a frame.
`close`	Finish the output. Returns the written path for `"usd"`.

Attributes:

Name	Type	Description
`frames`	`ndarray`	Captured frames as `(T, 3, H, W)` float32 -- the `"rtx"` path only.

Source code in xwm/envs/render.py

def __init__(
    self,
    model,
    *,
    backend: Backend = "usd",
    size: tuple[int, int] = (1280, 720),
    output_path: str | Path | None = None,
    environment: str = "studio",
    fps: int = 30,
    up_axis: str = "Z",
):
    if backend not in ("rtx", "usd"):
        raise ValueError(
            f"backend must be 'rtx' or 'usd' (got {backend!r}); for fast "
            "in-process frames use FrankaEnv.render, which is the 'warp' path"
        )
    requirement = RENDER_BACKENDS[backend]
    if not which_backends()[backend]:
        raise ImportError(
            f"the {backend!r} backend needs `pip install "
            f"{'ovrtx' if backend == 'rtx' else 'usd-core'}`"
            + (" and an NVIDIA GPU" if backend == "rtx" else "")
            + f" (missing module {requirement!r})"
        )

    import newton.viewer as viewer

    self.backend = backend
    self.size = size
    self.fps = fps
    self.up_axis = str(up_axis).upper()
    self._time = 0.0
    self._frames: list[np.ndarray] = []

    if backend == "rtx":
        self._viewer = viewer.ViewerRTX(
            width=size[0],
            height=size[1],
            headless=True,
            up_axis=up_axis,
            environment=environment,
            async_rendering=False,
        )
    else:
        if output_path is None:
            raise ValueError("the 'usd' backend needs an output_path")
        self.output_path = Path(output_path)
        self.output_path.parent.mkdir(parents=True, exist_ok=True)
        self._viewer = viewer.ViewerUSD(
            output_path=str(self.output_path), fps=fps, up_axis=up_axis
        )
    self._viewer.set_model(model)

frames ¶

frames: ndarray

Captured frames as (T, 3, H, W) float32 -- the "rtx" path only.

Feeds xwm.plots.save_gif and xwm.plots.plot_frames directly, so a path-traced episode is written exactly like a Warp one.

set_camera ¶

set_camera(*args, **kwargs) -> None

Forwarded to the viewer, where supported.

Source code in xwm/envs/render.py

def set_camera(self, *args, **kwargs) -> None:
    """Forwarded to the viewer, where supported."""
    setter = getattr(self._viewer, "set_camera", None)
    if setter is not None:
        setter(*args, **kwargs)

look_at ¶

look_at(eye, target) -> None

Aim the camera from eye at target, both in world metres.

Without this the path-traced view points at the horizon while the training camera looks at the arm, and the two renders are not of the same scene. See look_at_angles for the conversion.

Source code in xwm/envs/render.py

def look_at(self, eye, target) -> None:
    """Aim the camera from ``eye`` at ``target``, both in world metres.

    Without this the path-traced view points at the horizon while the
    training camera looks at the arm, and the two renders are not of the
    same scene. See :func:`look_at_angles` for the conversion.
    """
    pitch, yaw = look_at_angles(eye, target, self.up_axis)
    eye = np.asarray(eye, dtype=np.float64).reshape(3)
    self.set_camera(tuple(float(v) for v in eye), float(pitch), float(yaw))

add ¶

add(state) -> None

Record one simulation state as a frame.

Source code in xwm/envs/render.py

def add(self, state) -> None:
    """Record one simulation state as a frame."""
    self._viewer.begin_frame(self._time)
    self._viewer.log_state(state)
    self._viewer.end_frame()
    self._time += 1.0 / self.fps
    if self.backend == "rtx":
        self._frames.append(self._capture())

close ¶

close() -> Path | None

Finish the output. Returns the written path for "usd".

Source code in xwm/envs/render.py

def close(self) -> Path | None:
    """Finish the output. Returns the written path for ``"usd"``."""
    self._viewer.close()
    return getattr(self, "output_path", None)

discrete_action_table ¶

discrete_action_table(action_dim: int, *, magnitude: float = 1.0, include_noop: bool = True) -> ndarray

(n_actions, action_dim) table of axis-aligned moves.

Parameters:

Name	Type	Description	Default
`action_dim`	`int`	number of continuous action dimensions.	required
`magnitude`	`float`	how far each move pushes its dimension.	`1.0`
`include_noop`	`bool`	add an all-zeros action. Worth keeping: without it the agent cannot choose to stay put, which on a reach task means it can never stop once it arrives.	`True`

Returns:

Type	Description
`ndarray`	`2 * action_dim (+ 1)` rows, each a continuous action vector.

Source code in xwm/envs/discretize.py

def discrete_action_table(
    action_dim: int,
    *,
    magnitude: float = 1.0,
    include_noop: bool = True,
) -> np.ndarray:
    """``(n_actions, action_dim)`` table of axis-aligned moves.

    Args:
        action_dim: number of continuous action dimensions.
        magnitude: how far each move pushes its dimension.
        include_noop: add an all-zeros action. Worth keeping: without it the
            agent cannot choose to stay put, which on a reach task means it can
            never stop once it arrives.

    Returns:
        ``2 * action_dim (+ 1)`` rows, each a continuous action vector.
    """
    if action_dim < 1:
        raise ValueError(f"action_dim must be positive, got {action_dim}")
    moves = []
    for axis in range(action_dim):
        for sign in (+1.0, -1.0):
            action = np.zeros((action_dim,), np.float32)
            action[axis] = sign * magnitude
            moves.append(action)
    if include_noop:
        moves.append(np.zeros((action_dim,), np.float32))
    return np.stack(moves)

franka_sequences ¶

franka_sequences(env: FrankaEnv, n_sequences: int, length: int, *, seed: int = 0, smoothness: float = 0.7, progress_every: int = 0) -> dict[str, ndarray]

Collect an action-labelled video dataset from env.

The returned dict matches what xwm.action.ActionWorldModel expects, so it is a drop-in replacement for xwm.data.sprite_sequences:

video: (n, length, 3, H, W)
action: (n, length - 1, 7) -- action[i, t] joins frames t and t + 1
joint_q: (n, length, 7) ground truth, for probes
tool: (n, length, 3) ground-truth hand position, for probes

Diverged rollouts (NaN from the solver) are discarded and resampled, so the dataset never contains a corrupt sequence.

Source code in xwm/envs/newton_franka.py

def franka_sequences(
    env: FrankaEnv,
    n_sequences: int,
    length: int,
    *,
    seed: int = 0,
    smoothness: float = 0.7,
    progress_every: int = 0,
) -> dict[str, np.ndarray]:
    """Collect an action-labelled video dataset from ``env``.

    The returned dict matches what :class:`xwm.action.ActionWorldModel` expects,
    so it is a drop-in replacement for :func:`xwm.data.sprite_sequences`:

    * ``video``: ``(n, length, 3, H, W)``
    * ``action``: ``(n, length - 1, 7)`` -- ``action[i, t]`` joins frames
      ``t`` and ``t + 1``
    * ``joint_q``: ``(n, length, 7)`` ground truth, for probes
    * ``tool``: ``(n, length, 3)`` ground-truth hand position, for probes

    Diverged rollouts (NaN from the solver) are discarded and resampled, so the
    dataset never contains a corrupt sequence.
    """
    rng = np.random.default_rng(seed)
    videos, actions, joints, tools = [], [], [], []
    attempts = 0
    max_attempts = 4 * n_sequences + 16

    while len(videos) < n_sequences and attempts < max_attempts:
        attempts += 1
        action = smooth_actions(rng, 1, length - 1, env.action_dim, smoothness=smoothness)[0]
        result = env.rollout(action, seed=int(rng.integers(0, 2**31 - 1)))
        if not (env.is_finite() and np.all(np.isfinite(result["video"]))):
            continue
        videos.append(result["video"])
        actions.append(action)
        joints.append(result["joint_q"])
        tools.append(result["tool"])
        if progress_every and len(videos) % progress_every == 0:
            print(f"  collected {len(videos)}/{n_sequences} sequences", flush=True)

    if len(videos) < n_sequences:
        raise RuntimeError(
            f"only {len(videos)}/{n_sequences} rollouts stayed finite after {attempts} attempts; "
            "raise substeps or joint_armature in FrankaConfig"
        )
    return {
        "video": np.stack(videos),
        "action": np.stack(actions),
        "joint_q": np.stack(joints),
        "tool": np.stack(tools),
    }

smooth_actions ¶

smooth_actions(rng: Generator, n_sequences: int, length: int, action_dim: int, *, smoothness: float = 0.7) -> ndarray

(n, length, action_dim) temporally correlated actions in [-1, 1].

White noise makes an arm jitter in place and go nowhere; an AR(1) process produces trajectories that actually sweep through the workspace.

Source code in xwm/envs/newton_franka.py

def smooth_actions(
    rng: np.random.Generator,
    n_sequences: int,
    length: int,
    action_dim: int,
    *,
    smoothness: float = 0.7,
) -> np.ndarray:
    """``(n, length, action_dim)`` temporally correlated actions in ``[-1, 1]``.

    White noise makes an arm jitter in place and go nowhere; an AR(1) process
    produces trajectories that actually sweep through the workspace.
    """
    noise = rng.uniform(-1.0, 1.0, size=(n_sequences, length, action_dim)).astype(np.float32)
    out = np.empty_like(noise)
    carry = noise[:, 0]
    for t in range(length):
        carry = smoothness * carry + (1.0 - smoothness) * noise[:, t]
        out[:, t] = carry
    return np.clip(out, -1.0, 1.0)

look_at_angles ¶

look_at_angles(eye, target, up_axis: str = 'Z') -> tuple[float, float]

(pitch, yaw) in degrees for a viewer camera at eye facing target.

Newton's viewer camera is parameterised by position, pitch and yaw rather than by a look-at target. For a Z-up scene its forward vector is (cos yaw cos pitch, sin yaw cos pitch, sin pitch), which inverts to pitch = asin(dz) and yaw = atan2(dy, dx); the other two up-axes permute which components play those roles. Deriving the angles beats guessing them -- a camera aimed at the horizon puts the robot in a handful of pixels, which is a figure of nothing.

Source code in xwm/envs/render.py

def look_at_angles(eye, target, up_axis: str = "Z") -> tuple[float, float]:
    """``(pitch, yaw)`` in degrees for a viewer camera at ``eye`` facing ``target``.

    Newton's viewer camera is parameterised by position, pitch and yaw rather
    than by a look-at target. For a Z-up scene its forward vector is
    ``(cos yaw cos pitch, sin yaw cos pitch, sin pitch)``, which inverts to
    ``pitch = asin(dz)`` and ``yaw = atan2(dy, dx)``; the other two up-axes
    permute which components play those roles. Deriving the angles beats
    guessing them -- a camera aimed at the horizon puts the robot in a handful
    of pixels, which is a figure of nothing.
    """
    eye = np.asarray(eye, dtype=np.float64).reshape(3)
    direction = np.asarray(target, dtype=np.float64).reshape(3) - eye
    norm = float(np.linalg.norm(direction))
    if norm == 0.0:
        raise ValueError("camera position and target coincide")
    direction /= norm
    vertical = {"X": 0, "Y": 1, "Z": 2}[str(up_axis).upper()]
    # The two ground-plane axes, in the order the viewer's yaw sweeps them.
    first, second = {0: (1, 2), 1: (0, 2), 2: (0, 1)}[vertical]
    pitch = float(np.degrees(np.arcsin(np.clip(direction[vertical], -1.0, 1.0))))
    yaw = float(np.degrees(np.arctan2(direction[second], direction[first])))
    return pitch, yaw

supersample ¶

supersample(frame: ndarray, factor: int) -> ndarray

Box-downsample a (3, H*f, W*f) frame by factor.

Rendering above the target resolution and averaging down is the cheapest anti-aliasing there is, and the Warp raytracer has no built-in multisampling: one ray per pixel means every silhouette is a hard staircase. At factor=3 each output pixel integrates nine rays, which is enough to make edges and shadow boundaries read as smooth.

Source code in xwm/envs/render.py

def supersample(frame: np.ndarray, factor: int) -> np.ndarray:
    """Box-downsample a ``(3, H*f, W*f)`` frame by ``factor``.

    Rendering above the target resolution and averaging down is the cheapest
    anti-aliasing there is, and the Warp raytracer has no built-in
    multisampling: one ray per pixel means every silhouette is a hard staircase.
    At ``factor=3`` each output pixel integrates nine rays, which is enough to
    make edges and shadow boundaries read as smooth.
    """
    if factor <= 1:
        return frame
    channels, height, width = frame.shape
    if height % factor or width % factor:
        raise ValueError(f"frame {height}x{width} is not divisible by factor {factor}")
    reshaped = frame.reshape(channels, height // factor, factor, width // factor, factor)
    return reshaped.mean(axis=(2, 4))

which_backends ¶

which_backends() -> dict[str, bool]

Which rendering backends are usable here.

Worth calling before a long run: discovering that the high-quality path is unavailable after four hours of simulation is a poor use of a GPU.

Source code in xwm/envs/render.py

def which_backends() -> dict[str, bool]:
    """Which rendering backends are usable here.

    Worth calling before a long run: discovering that the high-quality path is
    unavailable after four hours of simulation is a poor use of a GPU.
    """
    import importlib.util

    available = {}
    for backend, requirement in RENDER_BACKENDS.items():
        available[backend] = requirement is None or (
            importlib.util.find_spec(requirement) is not None
        )
    return available