Franka manipulation
Train a Franka Panda arm to move its hand to a target that changes every episode. The policy emits only x, y, and z movement; a differential IK controller expands those three values into commands for the arm’s seven joints.
Save this complete app as franka_reach/app.py:
Show complete codefranka_reach/app.py
from __future__ import annotations
from typing import Any, Tuple
import simulo
# The Franka Panda arm — a global-catalog asset. Captured at module level# (construction-capture idiom): submit resolves and pins this exact version# before the job ever runs.franka = simulo.Asset.from_registry("simulo/robot/franka-panda:v1")
# A named, durable, writable volume for the trained checkpoint.vol = simulo.Volume.from_name("franka-reach-checkpoints", create_if_missing=True)
app = simulo.App("franka-reach", mounts={"/out": vol})
# The ONE module-level heavy import — deferred under the runtime guard.with app.runtime.imports(): import torch # noqa: F401 (resolved only in execution mode, on the worker)
@app.runtime.torch_jitdef _compute_rewards( rew_scale_distance: float, rew_scale_fine: float, rew_scale_action: float, distance_std: float, fine_std: float, ee_pos: torch.Tensor, goal_pos: torch.Tensor, actions: torch.Tensor,) -> torch.Tensor: """JIT-compiled reward kernel for task-space reaching.
Two nested distance terms plus an action penalty:
* a wide ``tanh`` shell that pulls the hand across the workspace toward the goal from anywhere, * a narrow one that only pays out in the last few centimetres, so the policy keeps improving after the wide term has saturated, * a small penalty on action magnitude, which stops the policy from thrashing the target point around once it is already on the goal. """ distance = torch.norm(goal_pos - ee_pos, dim=-1) rew_coarse = rew_scale_distance * (1.0 - torch.tanh(distance / distance_std)) rew_fine = rew_scale_fine * (1.0 - torch.tanh(distance / fine_std)) rew_action = rew_scale_action * torch.sum(torch.square(actions), dim=-1)
reward: torch.Tensor = rew_coarse + rew_fine + rew_action # Keep the (num_envs,) per-env reward contract even when num_envs == 1. return reward.view(-1)
class FrankaReachTask(simulo.Task): """Move a Franka Panda's hand onto a goal point that moves every episode.
Observation (23-dim): 7 joint positions (relative to the arm's rest pose), 7 joint velocities, the hand position, the goal position, and the vector from hand to goal — all positions in the arm's own base frame.
Action (3-dim): a Cartesian delta applied to the IK target point. The controller turns that into joint commands.
Defined at module level: ``simulo.Task`` is a torch-free contract stand-in at submit and the real training base on the worker, so the same class authors lean and trains heavy. """
observation_dim = 23 action_dim = 3
episode_length_s = 4.0
# How far one full-scale action moves the IK target point, per step [m]. action_scale = 0.05
# The link the controller drives and the joints it is allowed to move. # `panda_hand` is the wrist plate; the two finger joints are deliberately # NOT in this list — this task has no gripper. end_effector = "panda_hand" arm_joint_pattern = "panda_joint.*"
# Goal sampling box, in the arm's base frame [m]. Sized to sit inside the # Panda's comfortable reach so the IK solver always has an answer. goal_x_range = (0.35, 0.60) goal_y_range = (-0.25, 0.25) goal_z_range = (0.20, 0.50)
# Where the commanded point starts every episode — the centre of the goal # box, in the arm's base frame [m]. Fixed and identical for every episode, # so the policy's job (walk the point from here onto the goal) is a real # one and the first IK request of an episode is always a modest, bounded # move rather than a jump from wherever the last episode ended. initial_target = (0.475, 0.0, 0.35)
# The IK target is clamped to this box so a run of large actions cannot # walk the commanded point off into a region with no solution. target_x_range = (0.25, 0.70) target_y_range = (-0.40, 0.40) target_z_range = (0.10, 0.65)
# Hand orientation held for the whole episode: pointing straight down, # as a w-first quaternion. Reaching is a position task; pinning the # orientation keeps the arm in a sane posture without adding 3 more # action dimensions. ee_orientation = (0.0, 1.0, 0.0, 0.0)
rew_scale_distance = 1.0 rew_scale_fine = 0.5 rew_scale_action = -0.01 distance_std = 0.20 fine_std = 0.04
# Framework-injected at runtime by the training base (declared here only so # the type checker sees the names the methods read; the annotations are # PEP 563 strings and never shadow the inherited values). device: str num_envs: int max_episode_length: int episode_length_buf: torch.Tensor reset_terminated: torch.Tensor
def build(self, scene: simulo.Scene) -> None: scene.add(simulo.Terrain.plane(name="ground"), at="/", per_environment=False) scene.add( simulo.Light.dome(name="light", intensity=2500.0, color=(0.75, 0.75, 0.75)), at="/", per_environment=False, ) self.robot = simulo.Robot(asset=franka, initial_pose=simulo.Pose.identity()) scene.add(self.robot, at="/World/Robot")
def on_start(self, env: simulo.LearningEnv) -> None: self._arm_dof_idx = self.robot.find_joints(self.arm_joint_pattern)
# The controller is constructed against the ALREADY-BUILT robot, here # in on_start — it reads the robot's joints and bodies, which do not # exist during build(). `command_type="pose"` makes its command a # 7-vector: the hand's target position in the base frame plus the # w-first orientation quaternion held below. self.ik = simulo.DifferentialIKController( robot=self.robot, end_effector=self.end_effector, joints=self._arm_dof_idx, ik_method="dls", command_type="pose", )
# robot.state has no default-joint-value equivalent, so this stays on # the internals escape hatch (there is nothing unstable about reading # it here, just no supported, typed name for it yet). self._default_joint_pos = self.robot.internals.default_joint_pos.clone()
zeros = torch.zeros(self.num_envs, 3, device=self.device) self._goal_pos = zeros.clone() self._ee_pos = zeros.clone() self._actions = torch.zeros(self.num_envs, self.action_dim, device=self.device) self._orientation = torch.tensor(self.ee_orientation, device=self.device).repeat(self.num_envs, 1) self._initial_target = torch.tensor(self.initial_target, device=self.device).repeat(self.num_envs, 1) self._target_pos = self._initial_target.clone()
self._sample_goals(torch.arange(self.num_envs, device=self.device))
# -- helpers ----------------------------------------------------------
def _sample_goals(self, env_ids: torch.Tensor) -> None: """Draw a fresh goal point for each environment being reset.""" count = len(env_ids) for axis, (low, high) in enumerate((self.goal_x_range, self.goal_y_range, self.goal_z_range)): self._goal_pos[env_ids, axis] = torch.empty(count, device=self.device).uniform_(low, high)
def _read_ee_position(self) -> torch.Tensor: """Hand position in the arm's base frame, shape ``(num_envs, 3)``.
``get_body_pose_in_base_frame`` is the supported, typed readback; it returns ``(position, orientation)`` and this task only needs the position. Before the runtime attaches the robot it returns ``(None, None)``, so the zeroes below are the honest pre-attach value. """ position, _ = self.robot.get_body_pose_in_base_frame(self.end_effector) if position is None: return torch.zeros(self.num_envs, 3, device=self.device) return position
def _clamp_target(self) -> None: for axis, (low, high) in enumerate((self.target_x_range, self.target_y_range, self.target_z_range)): self._target_pos[:, axis] = self._target_pos[:, axis].clamp(low, high)
# -- the Task contract ------------------------------------------------
def get_observations(self) -> torch.Tensor: joint_pos = self.robot.state.joint_positions[:, self._arm_dof_idx] joint_vel = self.robot.state.joint_velocities[:, self._arm_dof_idx] joint_pos_rel = joint_pos - self._default_joint_pos[:, self._arm_dof_idx] return torch.cat( ( joint_pos_rel, joint_vel, self._ee_pos, self._goal_pos, self._goal_pos - self._ee_pos, ), dim=-1, )
def get_rewards(self) -> torch.Tensor: return _compute_rewards( self.rew_scale_distance, self.rew_scale_fine, self.rew_scale_action, self.distance_std, self.fine_std, self._ee_pos, self._goal_pos, self._actions, )
def get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: # Refresh the hand readback once per step, here, before the reward and # the next observation both read it. self._ee_pos = self._read_ee_position() truncated = self.episode_length_buf >= self.max_episode_length - 1 terminated = torch.zeros_like(truncated) return terminated, truncated
def apply_actions(self, actions: torch.Tensor) -> None: # This ALIASES the trainer's own action tensor -- nothing along # `SkrlWrapper.step -> LearningEnv.step -> on_pre_physics_step` copies # defensively, and skrl stores that same tensor as the transition's # action after `step()` returns. So: read `self._actions`, never write # into it. An in-place write here (or anywhere downstream) lands in a # transition PPO is about to record against a real reward. self._actions = actions self._target_pos = self._target_pos + self.action_scale * actions self._clamp_target() # A (num_envs, 7) pose command: position + the fixed w-first # orientation. The controller computes the joint targets and writes # them to the robot itself. self.ik.move_to(torch.cat((self._target_pos, self._orientation), dim=-1))
def reset_idx(self, env_ids: torch.Tensor) -> None: if len(env_ids) == 0: return self.robot.reset(env_ids)
# Back to the arm's rest pose with a little joint noise, so every # episode starts from a slightly different posture. The noise goes on # the SEVEN ARM JOINTS only: the Franka's two finger joints are # prismatic with 0.04 m of total travel, and +/-0.05 m of noise placed # them outside their own limits at the start of an episode. They are # not part of this task either -- its action is a 3-dim Cartesian # nudge, and it declares no gripper. joint_pos = self.robot.internals.default_joint_pos[env_ids].clone() arm = self._arm_dof_idx joint_pos[:, arm] += torch.empty_like(joint_pos[:, arm]).uniform_(-0.05, 0.05) joint_vel = self.robot.internals.default_joint_vel[env_ids] self.robot.set_joint_state(joint_pos, velocities=joint_vel, env_ids=env_ids)
self.ik.reset() self._sample_goals(env_ids) self._target_pos[env_ids] = self._initial_target[env_ids]
# Refresh the hand readback here as well as in `get_dones`. `env.reset()` # runs `reset_idx` and then `get_observations` -- it never calls # `get_dones` -- so without this the FIRST observation of a run reports # the hand at the origin and a goal-relative vector measured from it. self._ee_pos = self._read_ee_position()
# `self._actions[env_ids] = 0.0` used to live here. It was dead for # this task's arithmetic (`_actions` is read only by `get_rewards`, # which runs earlier in the same `env.step`, and `apply_actions` # rebinds it before the next read) and actively harmful otherwise: it # wrote zeros into the trainer's own tensor. See `apply_actions`.
@app.job( # Tier 1: T4 GPU, 16 GB VRAM. Run `simulo systems` for the full four-tier catalog. system=simulo.SystemType.TIER_1, timeout=4 * 60 * 60, retries=2, callbacks=[simulo.callbacks.ResumableCheckpoint(every=50)],)def train_franka_reach(num_envs: int = 2048, max_iterations: int = 300) -> dict[str, Any]: """Train a task-space reaching policy with PPO and save the checkpoint.
Args: num_envs: Number of parallel environments to simulate. max_iterations: Number of PPO policy-update iterations.
Returns: A JSON-serialisable dict: the saved ``checkpoint`` path plus training ``stats`` (``iterations``, ``num_envs``).
Raises: ValueError: If ``num_envs`` or ``max_iterations`` is not a positive integer. Checked here, at submit-shaped argument boundaries, because the alternative is a GPU job that boots the engine and then fails somewhere less legible — the same guard ``franka_lift`` carries. """ if num_envs < 1: raise ValueError(f"num_envs must be a positive integer, got {num_envs}") if max_iterations < 1: raise ValueError(f"max_iterations must be a positive integer, got {max_iterations}")
env = simulo.LearningEnv( task=FrankaReachTask(), num_envs=num_envs, device="cuda", dt=1.0 / 120.0, physics_steps_per_action=2, env_spacing=2.5, headless=True, seed=42, ) trainer = simulo.RLTrainer(env=env, algorithm="PPO", device="cuda", seed=42)
stats = trainer.train(max_iterations=max_iterations)
checkpoint = f"{vol.path}/franka_reach_final.pt" trainer.save(checkpoint)
# Close the trainer before the environment so the RL library releases its # resources first. trainer.close() env.close()
return {"checkpoint": checkpoint, "num_envs": num_envs, **stats}Run the app
Section titled “Run the app”simulo run franka_reach/app.py --num-envs 2048 --max-iterations 300Training logs will show the hand-to-target reward improving. The completed result
names the saved checkpoint. For a quick launch check, use --num-envs 64 --max-iterations 2.
What to reuse
Section titled “What to reuse”- Make the action describe the task—move the hand—not every joint in the robot.
- Read the hand pose in the robot’s base frame so observations and goals share a stable coordinate system.
- Reset only the environments that finish, including a fresh target for each one.
See Scene, Robot & World for scene construction and Tasks & the RL Loop for the training lifecycle.
Pick up a block
Section titled “Pick up a block”Now add a two-finger gripper and a movable Prop. This task rewards reaching,
grasping, lifting, and moving the block toward its goal. Success is measured from
the block’s height—not from closed fingers—so an empty grasp cannot look solved.
Save this complete app as franka_lift/app.py:
Show complete codefranka_lift/app.py
from __future__ import annotations
import osimport shutilfrom typing import Any, Tuple
import simulo
# The Franka Panda arm — a global-catalog asset, pinned at submit.franka = simulo.Asset.from_registry("simulo/robot/franka-panda:v1")
# A named, durable, writable volume for the trained checkpoint.vol = simulo.Volume.from_name("franka-lift-checkpoints", create_if_missing=True)
app = simulo.App("franka-lift", mounts={"/out": vol})
# The ONE module-level heavy import — deferred under the runtime guard.with app.runtime.imports(): import torch # noqa: F401 (resolved only in execution mode, on the worker)
@app.runtime.torch_jitdef _compute_rewards( rew_scale_reach: float, rew_scale_grasp: float, rew_scale_lift: float, rew_scale_goal: float, rew_scale_success: float, rew_scale_action: float, rew_scale_futile_close: float, reach_std: float, goal_std: float, grasp_radius: float, lift_height: float, success_height: float, close_threshold: float, grasp_pos: torch.Tensor, object_pos: torch.Tensor, goal_pos: torch.Tensor, object_height: torch.Tensor, resting_height: float, actions: torch.Tensor, futile_close: torch.Tensor,) -> torch.Tensor: """JIT-compiled staged reward kernel for pick-and-lift.
Five stages, each one reachable by gradient from the one before it — which is what makes a long-horizon manipulation task learnable at all:
* **reach** — a ``tanh`` shell on the distance from the point between the fingertips to the block. Pays out everywhere, so it pulls the hand in from anywhere in the workspace. * **grasp** — a bonus for having the fingertips actually around the block (that distance inside ``grasp_radius``) *and* commanding the fingers shut. The proximity half is measured on the *block's* pose, never on finger state: a parallel gripper reports the same thing whether it closed on the block or on nothing. * **lift** — how far off the table the block is, as a fraction of ``lift_height``. Continuous, and that is the whole point: see below. * **goal** — once lifted, a ``tanh`` shell on the distance from block to goal point, so the policy keeps improving after the lift itself is solved. Gated on ``lifted`` so it cannot be farmed by hovering over a block still sitting on the table. * **success** — a large one-off payment for getting the block above ``success_height``. Paid on the single step the episode ends, because that is also the step :meth:`FrankaLiftTask.get_dones` terminates on.
Plus one term that is not a stage: **futile close** charges the policy for commanding the gripper shut on a step the gate refused. It is the only shaping term here that is *identically zero on the desired policy* — see ``rew_scale_futile_close`` for what that buys and what it cannot buy.
Why the last two look the way they do — the defect they close -------------------------------------------------------------
The previous version of this kernel made **lift a step function** (``5.0`` if above ``lift_height``, ``0`` below) and had no success term at all. Both halves of that were wrong, and the training runs measured it:
1. **A step function has no gradient.** The trained policy plateaued with the block ``0.0103 m`` off the table, and the step paid nothing until ``0.035 m``. Between those two numbers the reward was flat, so nothing told the policy that 2 cm was better than 1 cm. Compare the reach term, which *is* shaped — and which converged beautifully. 2. **Succeeding cost more than it paid.** The episode ENDS on success, and reach+grasp pay ~1.34/step for up to 300 steps. A policy that lifted the block at step 120 forfeited 180 steps of that — about 240 reward — and the old lift term only paid 5.0/step for the handful of steps before termination. So a completed lift had *negative* advantage, and PPO correctly learned to hover and grip instead.
This defect was once explained as "the entropy run found genuine lifts at iterations 250 and 450 and never reinforced them". **That reading was overturned and does not survive**: re-reading both run logs end to end, the no-entropy run cleared 0.12 m twice and the entropy run twice — the same rate — and one of the no-entropy events is at *iteration 0*, an untrained policy that cannot grasp anything. Those were the arm batting the block, not lifting it. The negative-advantage arithmetic above stands on its own; it never needed that anecdote. See ``train_franka_lift``, which says the same thing about entropy at the point the knob is actually set.
Shaping alone would not have fixed it — it would have made things worse in an interesting way. A per-step lift reward that stops paying when the episode terminates means the best strategy is to sit *just below* the success bar collecting it forever. The one-off ``success`` payment is what makes crossing the bar strictly better than parking under it. The arithmetic is worked through in ``rew_scale_success``'s comment.
Plus a small penalty on action magnitude, which stops the policy thrashing the commanded point around once it is already where it wants to be. """ reach_distance = torch.norm(object_pos - grasp_pos, dim=-1) rew_reach = rew_scale_reach * (1.0 - torch.tanh(reach_distance / reach_std))
# The grasp bonus requires BOTH being on the object and commanding the # fingers shut. Rewarding proximity alone (the first version of this # kernel) is a dead end that measured clearly: the policy collects the # bonus by hovering, never closes, and so never reaches the lift term — # 0% success over 300 iterations with reward plateaued near 115. Nothing # in the reward mentioned the gripper, so nothing taught the policy to # use it. closing = (actions[:, 3] > close_threshold).float() grasped = (reach_distance < grasp_radius).float() * closing rew_grasp = rew_scale_grasp * grasped
# Height above the table, floored at zero: a block pressed INTO the table # pays nothing rather than paying a negative, which would make squashing it # a way to lose reward the policy has to actively avoid. lift = torch.clamp(object_height - resting_height, min=0.0) # Shaped, not a step. Every millimetre off the table pays, saturating at # lift_height — so there is a continuous gradient across exactly the band # (0.010 m to 0.035 m) the old step function left flat, and where the # trained policy measurably got stuck. rew_lift = rew_scale_lift * torch.clamp(lift / lift_height, max=1.0)
lifted = (lift > lift_height).float() goal_distance = torch.norm(goal_pos - object_pos, dim=-1) rew_goal = rew_scale_goal * lifted * (1.0 - torch.tanh(goal_distance / goal_std))
# The terminal payment. `get_dones` terminates on this same condition, so # the environment resets on the next step and this is collected exactly # once per episode — a one-off bonus, not an income stream. rew_success = rew_scale_success * (lift > success_height).float()
rew_action = rew_scale_action * torch.sum(torch.square(actions), dim=-1)
# The price of asking for something the gate will refuse. `futile_close` is # 1.0 exactly on the steps where the policy commanded close AND the latch # was shut, computed in `apply_actions` against the latch that governed THAT # action — never re-derived here from a post-physics latch one step newer # than the refusal it is meant to describe. rew_futile_close = rew_scale_futile_close * futile_close
reward: torch.Tensor = rew_reach + rew_grasp + rew_lift + rew_goal + rew_success + rew_action + rew_futile_close # Keep the (num_envs,) per-env reward contract even when num_envs == 1. return reward.view(-1)
@app.runtime.torch_jitdef _gripper_gate_permits( reach_distance: torch.Tensor, permitted: torch.Tensor, close_radius: float, release_radius: float,) -> torch.Tensor: """Which environments the gripper is allowed to close in, this step.
A Schmitt trigger on the fingertip-to-block distance, one bit per environment:
* **not currently permitted** — permission is granted when the fingertips come inside ``close_radius`` (the same number as ``grasp_radius``, so the gate opens exactly where the reward says a grasp counts). * **currently permitted** — permission is kept until the fingertips leave the larger ``release_radius``.
Why two radii instead of one ----------------------------
A single threshold makes the gate chatter. A policy that converges to hovering *at* the radius — which is exactly what this one does, at 0.016–0.026 m against a 0.035 m radius — sits in the noise band, and the block's own contact jitter then flips permission on and off every control step. That would slam the fingers open and shut at 60 Hz, and a parallel gripper needs two consecutive settled updates before it will even report ``CLOSED``, so a chattering gate cannot produce a grasp at all.
The hysteresis band is also the reason the gate's state has to be in the observation: with one radius the gate would be a pure function of the block-to-fingertips vector the policy already sees, but a latch depends on how the environment *got* here. See ``FrankaLiftTask.get_observations``.
Args: reach_distance: Fingertip-point-to-block distance, ``(num_envs,)`` [m]. permitted: Last step's permission bit, ``(num_envs,)`` bool. close_radius: Distance inside which permission is granted [m]. release_radius: Distance outside which permission is withdrawn [m]. Must be the larger of the two, or the trigger inverts into a one-shot latch that never releases.
Returns: The new permission bit, ``(num_envs,)`` bool. """ return torch.where(permitted, reach_distance < release_radius, reach_distance < close_radius)
@app.runtime.torch_jitdef _gated_gripper_command( grip_command: torch.Tensor, permitted: torch.Tensor, open_command: float,) -> torch.Tensor: """The policy's gripper channel after the gate has had its say.
Inside the permitted band the policy's own value passes through untouched, so the gate is *invisible* exactly where a grasp is possible. Outside it, the channel is overwritten with ``open_command`` — not with "idle", and that is deliberate: idle holds the current target, which would let a gripper that closed and then lost permission stay shut around nothing. The invariant worth having is the strong one — **outside the band, the gripper is open.**
Args: grip_command: The policy's raw fourth action value, ``(num_envs,)``. permitted: Whether each environment may close, ``(num_envs,)`` bool. open_command: Value that means "open" on the gripper's command channel (``GripperCommand.OPEN``, i.e. -1.0 — anything below the actuator's ``open_threshold`` would do, but the enum's own value is the one that says why).
Returns: The gripper command actually written, ``(num_envs,)``. """ return torch.where(permitted, grip_command, torch.full_like(grip_command, open_command))
class FrankaLiftTask(simulo.Task): """Pick a block up off the table and carry it to a point above it.
Observation (42-dim): 7 joint positions (relative to the arm's rest pose), 7 joint velocities, the hand's position and orientation, the block's position and orientation, the goal position, the block-to-fingertips and block-to-goal vectors, the previous action, and the gripper gate's state — every position in the arm's own base frame.
Action (4-dim): a Cartesian delta applied to the IK target point (3), plus a continuous gripper command (1). The controller turns the first three into joint commands; the fourth goes to the gripper **through the gate** — see :meth:`apply_actions`, which is where this task stops being a pure pass-through of the policy's action.
Defined at module level: ``simulo.Task`` is a torch-free contract stand-in at submit and the real training base on the worker, so the same class authors lean and trains heavy. """
observation_dim = 42 action_dim = 4
episode_length_s = 5.0
# How far one full-scale action moves the IK target point, per step [m]. action_scale = 0.04
# The link the controller drives and the joints it is allowed to move. # The two finger joints are deliberately NOT in this list — the gripper # owns those, and nothing else may write to them. end_effector = "panda_hand" arm_joint_pattern = "panda_joint.*" finger_joint_pattern = "panda_finger_joint.*"
# Finger joint position for a fully open / fully closed gripper [m]. The # Panda's fingers travel 0 to 0.04 each, so 0.04 is an 0.08 m aperture. open_width = 0.04 closed_width = 0.0
# Distance from the ``panda_hand`` frame down to the point the fingers # close on [m]. The hand's orientation is pinned pointing straight down for # the whole episode, so this is a pure -z offset, and it is what turns # "where the hand frame is" into "where the fingers would close". # # 0.103 is the fingerTIP offset, measured on this arm: the finger bodies' # origins sit 0.058 below the hand frame and the pads run roughly 0.045 # further down. # # This offset is the difference between the two frames the task juggles, and # mixing them up is the easy mistake here: the IK controller commands the # HAND, while the reward measures the point the FINGERS close on. Command # the hand straight to the block's centre and the fingertips end up 0.103 m # below it — through the table. A scripted approach-grasp-lift expert that # adds this offset lifts the block in 100% of environments; the same script # without it lifts in 6%. hand_to_grasp_point = 0.103
# -- the table -------------------------------------------------------- #: Height of the work surface the block rests on [m]. The block sits on a #: table rather than on the floor so the fingers have somewhere to go: the #: gripper has to reach BELOW the block's centre to close around it, and on #: an infinite ground plane that is straight into the floor. #: #: Set to 0.0 to remove the table and drop the block on the ground — the #: control condition used to measure what the table is worth. table_height = 0.10 #: Footprint of the table top [m], centred under the block region. table_size = (0.60, 1.20) table_center_x = 0.55
# -- the block -------------------------------------------------------- #: Edge length of the block [m]. Narrower than the gripper's 0.08 m open #: aperture, so the fingers close around it rather than onto it. block_size = 0.05 #: Where the block spawns, in the world frame [m] — resting on the table #: (or on the ground when ``table_height`` is 0), so it starts settled #: rather than dropping. #: #: **This is the value for the DEFAULT ``table_height``, and nothing reads #: its z at run time** — see :meth:`block_resting_z`. The distinction is #: not pedantry: this line is evaluated once, when the class body is #: executed, so a subclass that sets ``table_height = 0.0`` (the control #: condition this demo documents) would otherwise keep a z computed from #: 0.10 and spawn the block 10 cm inside the floor. Every derived quantity #: is therefore computed from ``self`` at run time instead. block_position = (0.5, 0.0, table_height + block_size / 2)
#: How far the block's spawn point is perturbed at each reset [m]. This is #: the sampled region the block is re-randomised over, per environment. block_x_jitter = 0.10 block_y_jitter = 0.20
# -- the goal --------------------------------------------------------- # The goal point sits directly above the block's spawn, at # ``goal_height`` (defined with the reward parameters below, because it is # pinned to ``success_height``). It used to be sampled from an ABSOLUTE # band in the arm's base frame — see ``goal_height`` for the defect that # was, and why the band is gone.
# Where the commanded IK point starts every episode, in the arm's base # frame [m] — above the middle of the block region, so the first IK request # of an episode is always a modest, bounded move. Same caveat as # `block_position`: this is the DEFAULT, and `on_start` derives the live # value from `self` so it tracks `table_height`. initial_target = (0.5, 0.0, block_position[2] + hand_to_grasp_point + 0.10)
# The IK target is clamped to this box so a run of large actions cannot # walk the commanded point off into a region with no solution. target_x_range = (0.25, 0.70) target_y_range = (-0.40, 0.40) target_z_range = (0.06, 0.65)
# Hand orientation held for the whole episode: pointing straight down, as # a w-first quaternion. Pinning it keeps the arm in a sane grasping posture # without adding 3 more action dimensions. ee_orientation = (0.0, 1.0, 0.0, 0.0)
# -- reward shaping --------------------------------------------------- rew_scale_reach = 1.0 rew_scale_grasp = 0.5 #: Peak per-step payment for the shaped lift term, reached at #: ``lift_height``. Was ``5.0`` when the term was a step function; it is #: lower now precisely BECAUSE the term is continuous — a per-step reward #: that is collectable for all 300 steps of an episode compounds, and the #: bigger this is the bigger ``rew_scale_success`` has to be to outrank #: parking just below the success bar. Small keeps the spread sane. rew_scale_lift = 1.0 #: Peak per-step payment for the goal term — **zero, deliberately**, and #: that is a finding rather than a disabled knob. #: #: The goal stage's stated job is "once lifted, keep improving after the #: lift itself is solved". :meth:`get_dones` ends the episode the moment #: the block passes ``success_height``, so there is no "after" to improve #: in: the goal term is live only in the sliver between ``lift_height`` #: (0.035 m, where ``lifted`` switches on) and ``success_height`` #: (0.05 m, where the episode stops). A 1.5 cm window is not a carry #: stage. #: #: Paying it anyway costs something real, which is why this is 0.0 and not #: a small number. The term is per-step and gated on ``lifted``, so it is #: collected up to 300 times by a policy that parks just below the bar and #: exactly once by a policy that crosses it. Any live goal term therefore #: erodes the success-over-parking margin the ``rew_scale_success`` #: arithmetic was sized to guarantee. The concrete returns below make the #: comparison explicit: with the goal term off the margin is **2.48x**; at #: ``rew_scale_goal = 0.5`` it is **2.07x**, and at 1.0 it is **1.77x**, #: because the parking return climbs 701 -> 851 -> 1001 while the success #: return barely moves (1740 -> 1758 -> 1775). Restoring the term would mean #: raising ``rew_scale_success``, i.e. re-opening the reward rebalance. #: #: (This note used to say the margin fell "from 2.4x to about 1.2x". The #: first figure is right; **1.2x was not a measurement of anything** — no #: setting of this knob produces it, and running the shipped helpers is how #: that was established.) #: #: So this task pays FOUR of the kernel's five stages. The fifth belongs to #: a variant that terminates on reaching the goal rather than on clearing a #: height, which is a different task. rew_scale_goal = 0.0 #: One-off payment for ending the episode above ``success_height``. #: #: Sized from the returns it has to beat, not picked for feel. Measured #: per-step values at convergence: reach ~0.844 (at a 0.016 m reach error), #: grasp 0.5, shaped lift up to 1.0, goal up to ~0.46 while lifted. #: #: * hover on the table, grip, never lift — what the policy actually #: learned: ~295 over 300 steps (measured, three runs). #: * park just below the bar, collecting the shaped lift forever — the #: local optimum shaping ALONE would create: ~700. #: * grasp, lift, succeed at ~step 120: ~204 of accumulated shaping, plus #: this bonus. #: #: At 1500 that last line is ~1704 — 2.4x the parking strategy and 5.8x #: hovering, with a gradient running from each to the next. The ordering is #: the requirement; the exact number just needs enough margin that it #: survives the reach term being worth a bit more or less than assumed. rew_scale_success = 1500.0 rew_scale_action = -0.005 #: Per-step charge for commanding the gripper shut on a step the gate #: refused — see :meth:`apply_actions`, which computes the bit. #: #: **What this is for.** Outside the latch the MDP is otherwise exactly #: invariant to action channel 3: ``apply_actions`` overwrites it before it #: reaches the actuator, and ``_compute_rewards`` multiplies ``closing`` by #: ``reach_distance < grasp_radius``. So the policy is not stubborn about #: closing early, it is **indifferent** — nothing it does on that channel in #: the far field changes anything it is paid. The measured consequence is #: the commanded-versus-effective table above: close is still commanded at #: step 1.23 from 0.252 m, in 95.6% of first closes outside the radius. #: #: **Why a price rather than more shaping.** This term is identically zero #: on the desired policy — a policy that only closes inside the radius never #: pays it — so it cannot move the optimum, only break the tie in the region #: where the optimum is currently flat. Every previous round of shaping on #: this task removed one degenerate optimum by installing another; a term #: that is zero at the optimum cannot do that. #: #: **The sizing, and the three numbers that bracket it.** #: #: * **20x the 0.005 that measurably failed.** ``rew_scale_action`` is #: computed on the RAW action, so a gradient toward channel-3 = 0 in far #: states already existed at 0.005/step and demonstrably did not move the #: behaviour. That is an empirical lower bound: 0.005/step is below PPO's #: effective noise floor here, where skrl normalises GAE advantages and #: variance is dominated by the 1500-point terminal bonus plus the #: documented non-monotonic collapses. #: * **5x under the 0.5 grasp bonus**, so it can never out-price closing #: near the block. Inside the latch it is not merely outranked, it is #: zero. #: * **Worst case 300 x 0.1 = 30** over a full episode, against a strategy #: table of hover ~295 / park ~700 / succeed ~1704. Only the *ordering* of #: those is load-bearing (see ``rew_scale_success``), and 30 cannot #: reorder them. #: #: **The degenerate optimum it installs, stated because it is real.** #: "Never close" — during discovery rather than at equilibrium. Early on #: nearly every close is futile, so the initial gradient is "lower channel 3 #: everywhere"; the exception ("except near the block") is only carved back #: out if the policy still samples close inside the latch often enough. If #: global suppression wins that race, the run re-converges to reach-and- #: hover at ~0%. The instrumentation for it is peak lift, the mean channel-3 #: output split by gate state, and the per-step grasp-bonus incidence. rew_scale_futile_close = -0.1 reach_std = 0.10 goal_std = 0.20 #: How close the fingertip point must be to the block's centre to count as #: a grasp [m]. Roughly half the block's diagonal — the fingers are around #: it at this range. grasp_radius = 0.035
# -- the gripper gate ------------------------------------------------- #: Fingertip-to-block distance inside which the gate lets the policy close #: the gripper [m]. #: #: **This is ``grasp_radius``, and the identity is the whole design.** The #: gate is transparent exactly where the reward pays the grasp bonus, so it #: cannot change the reward's value anywhere: the only place ``closing`` #: matters to ``_compute_rewards`` is inside ``grasp_radius``, which is #: precisely where the gate passes the policy's command through untouched. #: That is what makes this an ACTION gate and not another round of reward #: shaping: inside the radius the raw command passes; outside it, the #: gripper is held open. gate_close_radius = grasp_radius #: Distance at which the gate takes permission away again [m]. #: #: Larger than ``gate_close_radius`` on purpose — the two numbers are a #: Schmitt trigger, see ``_gripper_gate_permits``. 0.06 m is comfortably #: past any in-hand jitter of a block being carried (the fingertip point #: tracks a held block, so the distance stays near zero) and comfortably #: short of "the block is gone". gate_release_radius = 0.06 #: Command value above which the gripper closes. Must match the #: ``close_threshold`` the ParallelGripperActuator is built with (its #: default), because the reward reads the policy's raw gripper action and #: has to agree with the actuator about what counts as "closing". gripper_close_threshold = 0.3 #: How far off the table the block must be raised to earn the lift term [m]. lift_height = 0.035 #: How far off the table the block must be raised to END the episode #: successfully [m]. #: #: 0.05 m is "the block is clearly off the table", which is what a lift #: demo should mean by success. It was 0.12 m, and that was mis-calibrated: #: a *perfect scripted expert* on this scene peaks at about **0.126 m**, so #: the bar sat within 5% of the best an optimal open-loop controller can #: do. That measures "lift as well as a flawless script", not "lift the #: block", and it makes a partly-working policy indistinguishable from one #: that never moves the block at all. #: #: Recalibrating did not flatter the numbers at the time it was done: the #: then-current policy's best peak lift was 0.0122 m, which scores 0% #: against either bar. The threshold was wrong independently of what it #: happened to score, which is why it was worth fixing rather than keeping #: as a convenient excuse. (The current policy's peak lift at iteration 1150 #: averages 0.0723 m and tops out at 0.0811 m across 2048 environments, so #: it clears this bar comfortably and would still miss the old one — which #: is exactly the range of outcomes a 0.12 m bar could not distinguish. The #: 0.079 m this note used to quote was cycle FIVE's population maximum: one #: cycle stale, and a different statistic from the table's mean.) success_height = 0.05
#: How far above the block's RESTING height the goal point sits [m]. #: #: Two things about this line are the fix for a real defect, and both are #: worth stating because the broken version looked perfectly reasonable. #: #: **It is ``success_height``, not an independent number.** The goal used #: to be drawn from ``goal_z_range = (0.25, 0.40)`` — an *absolute* band in #: the arm's base frame. :meth:`get_dones` terminates the episode at #: ``resting_height + success_height`` = 0.175 m, so every goal in that #: band sat above the height at which the episode stops existing. The goal #: was unreachable by construction, and no amount of training could have #: found it. Tying the two together is what makes that class of mismatch #: unrepresentable: the goal IS the bar. #: #: **It is measured from the resting height, not from the world origin.** #: The old band was absolute, so it silently stopped meaning anything when #: ``table_height`` changed — and ``table_height = 0.0`` is the documented #: control condition this demo tells you to run. A block resting at 0.025 m #: with a goal at 0.25–0.40 m is a 22–37 cm lift, against a success bar of #: 7.5 cm. Relative placement tracks the table automatically. goal_height = success_height
# Framework-injected at runtime by the training base (declared here only so # the type checker sees the names the methods read; the annotations are # PEP 563 strings and never shadow the inherited values). device: str num_envs: int max_episode_length: int episode_length_buf: torch.Tensor reset_terminated: torch.Tensor
def block_resting_z(self) -> float: """The block's settled centre height [m], derived at RUN time.
Everything that needs "where the block sits" goes through here rather than through ``block_position[2]``, because a class-body expression is evaluated once and a subclass cannot change it afterwards. ``table_height = 0.0`` is a control condition this demo tells you to run, and setting it by subclassing used to leave the spawn height, the lift datum, the goal height and the success bar all computed from a table that is no longer there — a silent 0%, with no error and no obviously wrong number to notice. """ return float(self.table_height) + float(self.block_size) / 2.0
def initial_target_position(self) -> Tuple[float, float, float]: """The episode's starting IK target [m], derived at RUN time.
The x and y come from ``initial_target``; only the z is recomputed, for the reason :meth:`block_resting_z` gives. """ return ( float(self.initial_target[0]), float(self.initial_target[1]), self.block_resting_z() + self.hand_to_grasp_point + 0.10, )
def build(self, scene: simulo.Scene) -> None: scene.add(simulo.Terrain.plane(name="ground"), at="/", per_environment=False) scene.add( simulo.Light.dome(name="light", intensity=2500.0, color=(0.75, 0.75, 0.75)), at="/", per_environment=False, ) self.robot = simulo.Robot(asset=franka, initial_pose=simulo.Pose.identity()) scene.add(self.robot, at="/World/Robot")
# The table. A plain Cuboid, NOT a Prop — nothing ever moves it, so it # needs no pose surface. `kinematic=True` is what makes it a fixed # collider: it collides and is never pushed, which is exactly a table. if self.table_height > 0.0: scene.add( simulo.Cuboid( name="Table", size=(self.table_size[0], self.table_size[1], self.table_height), pose=simulo.Pose(position=[self.table_center_x, 0.0, self.table_height / 2.0]), material=simulo.Material.surface(color=(0.35, 0.30, 0.26)), physics=simulo.Physics.rigid(kinematic=True, collision=simulo.Collision.enabled()), ), at="/World/Table", )
# The block, as a Prop rather than a bare Cuboid. The shape is exactly # what it always was — the Prop wrapper is what gives the task a pose # it can read back and write, which is the whole reason this app can be # trained instead of scripted. self.block = simulo.Prop( simulo.Cuboid( name="Block", size=(self.block_size, self.block_size, self.block_size), # z from `block_resting_z()`, not `block_position[2]`, so the # block still rests ON the surface when a subclass moves the # table (including removing it). pose=simulo.Pose(position=[self.block_position[0], self.block_position[1], self.block_resting_z()]), material=simulo.Material.surface(color=(0.9, 0.35, 0.15)), physics=simulo.Physics.rigid(mass=0.05), ) ) scene.add(self.block, at="/World/Block")
def on_start(self, env: simulo.LearningEnv) -> None: self._arm_dof_idx = self.robot.find_joints(self.arm_joint_pattern)
# Both are constructed against the ALREADY-BUILT robot, here in # on_start — each reads joints and bodies that do not exist during # build(). They drive DISJOINT joint sets: the controller owns the # seven arm joints, the gripper owns the two finger joints. self.ik = simulo.DifferentialIKController( robot=self.robot, end_effector=self.end_effector, joints=self._arm_dof_idx, ik_method="dls", command_type="pose", ) self.gripper = simulo.ParallelGripperActuator( robot=self.robot, finger_joints=self.finger_joint_pattern, open_width=self.open_width, closed_width=self.closed_width, )
# robot.state has no default-joint-value equivalent, so this stays on # the internals escape hatch (there is nothing unstable about reading # it here, just no supported, typed name for it yet). self._default_joint_pos = self.robot.internals.default_joint_pos.clone()
self._goal_pos = torch.zeros(self.num_envs, 3, device=self.device) self._grasp_pos = torch.zeros(self.num_envs, 3, device=self.device) self._ee_quat = torch.zeros(self.num_envs, 4, device=self.device) self._object_pos = torch.zeros(self.num_envs, 3, device=self.device) self._object_quat = torch.zeros(self.num_envs, 4, device=self.device) self._actions = torch.zeros(self.num_envs, self.action_dim, device=self.device)
# The gate's latch: may this environment close the gripper? Starts # shut, so every episode begins with an approach the policy cannot # slam the fingers through. self._gate_open = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) # "Did the policy ask to close on a step the gate refused?" — written by # `apply_actions`, read by `get_rewards`, one bit per environment. # Allocated here so the first `get_rewards` of a run has something to # read even if it somehow precedes an `apply_actions`. self._futile_close = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) # The value written to the gripper channel while the gate is shut. # Hoisted here rather than inlined as -1.0 because `on_start` runs only # on the worker, where `simulo.GripperCommand` is the real enum rather # than the torch-free submit stand-in — so the demo can say what it # means with the typed vocabulary instead of a magic number. self._gate_shut_command = float(simulo.GripperCommand.OPEN)
self._orientation = torch.tensor(self.ee_orientation, device=self.device).repeat(self.num_envs, 1) self._initial_target = torch.tensor(self.initial_target_position(), device=self.device).repeat(self.num_envs, 1) self._target_pos = self._initial_target.clone() self._grasp_offset = torch.tensor([0.0, 0.0, self.hand_to_grasp_point], device=self.device)
# The block's resting height in the base frame — the datum every lift # is measured against. Derived here rather than read off # `block_position`, so it tracks `table_height`. self._resting_height = self.block_resting_z() # Passed to gripper.update() each step. A parallel gripper ignores dt # (it infers its state from the finger joints, which the engine has # already advanced), but pass the real value rather than a placeholder. self._gripper_dt = float(getattr(env, "dt", 1.0 / 120.0))
all_envs = torch.arange(self.num_envs, device=self.device) self._refresh_readbacks() self._update_gripper_gate() self._place_goals(all_envs, self._object_pos)
# -- helpers ----------------------------------------------------------
def _base_frame_offset(self) -> torch.Tensor: """The arm's base position in the world frame, ``(num_envs, 3)``.
The block's pose comes back in the WORLD frame while every other quantity this task uses is in the arm's own base frame, so one of them has to be converted. Subtracting the arm's root position is the whole conversion here — and only because this app spawns the arm at ``Pose.identity()``, so the base frame is the world frame translated, never rotated. An arm spawned at an angle would need the rotation too. """ return self.robot.state.pose[:, :3]
def _place_goals(self, env_ids: torch.Tensor, block_pos_b: torch.Tensor) -> None: """Put the lift goal above the block for each environment being reset.
The goal sits directly above wherever the block was just placed, at ``goal_height`` — so it names a clean vertical lift rather than a drag sideways. It is a fixed offset, not a sampled one: see ``goal_height`` for why the height it used to be drawn from was unreachable.
Even though ``rew_scale_goal`` is 0.0, this is not dead code — the goal point and the block-to-goal vector are six of the observation's dimensions, so the policy is told the height it has to clear. That is strictly more useful than the random unreachable height it used to be told.
Args: env_ids: The environments being reset. block_pos_b: Where the block now is for those environments, in the arm's base frame. Passed in rather than read back on purpose — see the caller in :meth:`reset_idx`. """ self._goal_pos[env_ids, 0] = block_pos_b[:, 0] self._goal_pos[env_ids, 1] = block_pos_b[:, 1] self._goal_pos[env_ids, 2] = self._resting_height + self.goal_height
def _refresh_readbacks(self) -> None: """Read the hand and the block once per step, before anything uses them.
``get_body_pose_in_base_frame`` is the supported, typed hand readback; ``block.state`` is the supported, typed block readback. Both are re-read every step rather than held, which is the documented contract for live state views. """ ee_pos, ee_quat = self.robot.get_body_pose_in_base_frame(self.end_effector) if ee_pos is not None: # The point the fingers would close on, not the hand frame itself. self._grasp_pos = ee_pos - self._grasp_offset self._ee_quat = ee_quat
block_pose = self.block.state.pose self._object_pos = block_pose[:, :3] - self._base_frame_offset() self._object_quat = block_pose[:, 3:7]
def _update_gripper_gate(self) -> None: """Recompute the gate's latch from the readbacks just refreshed.
Called from :meth:`on_post_physics_step`, and that placement is load bearing rather than incidental. The environment's step order is
apply_actions -> physics -> on_post_physics_step -> dones -> rewards -> observations
so updating the latch here means the bit :meth:`get_observations` reports is exactly the bit :meth:`apply_actions` will apply to the next action. The policy is never told one thing and given another. """ reach_distance = torch.norm(self._object_pos - self._grasp_pos, dim=-1) self._gate_open = _gripper_gate_permits( reach_distance, self._gate_open, self.gate_close_radius, self.gate_release_radius, )
def _clamp_target(self) -> None: for axis, (low, high) in enumerate((self.target_x_range, self.target_y_range, self.target_z_range)): self._target_pos[:, axis] = self._target_pos[:, axis].clamp(low, high)
# -- the Task contract ------------------------------------------------
def on_post_physics_step(self) -> None: # Refresh both readbacks once per step, here — this hook runs before # get_dones, get_rewards, and get_observations, all three of which read # them. self._refresh_readbacks() # The gripper's reported state is refreshed ONLY by update(): without # this call, get_state() and is_closed() keep returning the value the # actuator was constructed with (OPEN) for the whole run, silently. This # task does not read gripper state in its reward — success is measured # on the block — but leaving it unwound makes the actuator lie to # anything that does look, including debugging. self.gripper.update(self._gripper_dt) # Last of the three, because it reads the readbacks refreshed above. self._update_gripper_gate()
def get_observations(self) -> torch.Tensor: """The 42-dim observation, including one bit for the gripper gate.
Why the gate's state is in here — the decision, and the reasoning --------------------------------------------------------------------
Gating the gripper makes the *effect* of action 3 depend on something other than action 3. Left invisible, that is a non-stationary action effect, which is a learning problem in its own right: the same command does two different things and nothing in the observation says which.
The tempting counter-argument is that the gate is redundant — the observation already carries ``object_pos - grasp_pos``, and the gate is a threshold on that vector's norm, so the policy could compute it. That argument is **correct for a single-threshold gate and wrong for this one.** The hysteresis band (see ``_gripper_gate_permits``) makes permission depend on how the environment arrived: at a distance of 0.05 m the gate is open if the fingertips came from inside 0.035 m and shut if they did not, and nothing in an observation of the *current* state distinguishes those. Adding hysteresis is what creates the hidden state, so adding hysteresis is what obliges the bit.
The cost is one dimension out of 42. That is not a close call.
The bit is the latch that will govern the NEXT action, not a report of what happened to the last one — see :meth:`_update_gripper_gate` for why the ordering works out that way. """ joint_pos = self.robot.state.joint_positions[:, self._arm_dof_idx] joint_vel = self.robot.state.joint_velocities[:, self._arm_dof_idx] joint_pos_rel = joint_pos - self._default_joint_pos[:, self._arm_dof_idx] return torch.cat( ( joint_pos_rel, joint_vel, self._grasp_pos, self._ee_quat, self._object_pos, self._object_quat, self._goal_pos, self._object_pos - self._grasp_pos, self._goal_pos - self._object_pos, self._actions, self._gate_open.unsqueeze(-1).float(), ), dim=-1, )
def get_rewards(self) -> torch.Tensor: return _compute_rewards( self.rew_scale_reach, self.rew_scale_grasp, self.rew_scale_lift, self.rew_scale_goal, self.rew_scale_success, self.rew_scale_action, self.rew_scale_futile_close, self.reach_std, self.goal_std, self.grasp_radius, self.lift_height, self.success_height, self.gripper_close_threshold, self._grasp_pos, self._object_pos, self._goal_pos, self._object_pos[:, 2], self._resting_height, self._actions, self._futile_close.float(), )
def get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]: truncated = self.episode_length_buf >= self.max_episode_length - 1 # Success is measured on the BLOCK: it is off the table, which nothing # but a real grasp can achieve. Finger state is deliberately not # consulted — a parallel gripper reports CLOSED whether it closed on # the block or on empty air. terminated = self._object_pos[:, 2] > self._resting_height + self.success_height return terminated, truncated
def apply_actions(self, actions: torch.Tensor) -> None: """Turn the policy's four numbers into an IK command and a gripper command.
This is where the gate lives, and the placement is the design decision rather than a convenience. The gate is a statement about *this task's* action semantics — "you may not close the fingers until they are around the block" — and ``Task.apply_actions`` is the one place that owns the mapping from action vector to actuator commands.
The alternative was to put it in ``ParallelGripperActuator``. That would have been wrong twice over: the actuator is a general SDK type that knows about finger joints and nothing about blocks or ``grasp_radius``, so the gate would have had to be handed the very task-specific state that makes it task-specific; and a reusable gripper that silently refuses commands is a worse gripper. Keeping it here leaves the SDK honest and leaves the demo's one unusual rule visible in the demo. """ # Clone: the trainer reuses and later reads this tensor, so keeping a # bare alias here and writing through it would corrupt the action the # algorithm believes it took. self._actions = actions.clone()
# Did the policy ask for something the gate is about to refuse? Computed # HERE, before the gate is applied, because `self._gate_open` at this # instant is the latch that governs THIS action. Recomputing it in the # reward would read the latch `on_post_physics_step` has since updated — # one step newer than the refusal it is supposed to describe. self._futile_close = (self._actions[:, 3] > self.gripper_close_threshold) & ~self._gate_open
self._target_pos = self._target_pos + self.action_scale * self._actions[:, :3] self._clamp_target() # A (num_envs, 7) pose command: position + the fixed w-first # orientation. The controller computes the joint targets and writes # them to the robot itself. self.ik.move_to(torch.cat((self._target_pos, self._orientation), dim=-1))
# The gripper's own channel: the policy's continuous fourth value, one # per environment — passed through the gate rather than straight # through. Inside `gate_close_radius` the gate is invisible and the # dead band still does its job of letting a policy HOLD a grasp; # outside it the channel reads OPEN whatever the policy asked for. # # `self._actions` deliberately keeps the RAW command, not the gated # one: it is what the action penalty is computed on and what the # observation reports as "previous action", and both should describe # what the policy chose. What the gate did about it is its own # observation bit. # # Keeping it raw is also what makes the grasp bonus and the futile-close # penalty two prices on the SAME quantity — the command the policy # chose. Read the bonus off the raw channel and the penalty off the # gated one and they would be pricing different things, which is how a # shaping term stops being a tie-breaker and starts being an argument. self.gripper.set_commands(_gated_gripper_command(self._actions[:, 3], self._gate_open, self._gate_shut_command))
def reset_idx(self, env_ids: torch.Tensor) -> None: if len(env_ids) == 0: return self.robot.reset(env_ids)
# Back to the arm's rest pose with a little joint noise, so every # episode starts from a slightly different posture. The noise goes on # the SEVEN ARM JOINTS only, and on this task that restriction is not # tidiness — it is a correctness fix, for the same reason # `franka_reach` carries it. # # The Franka's two finger joints are prismatic with 0.04 m of travel # each, so +/-0.05 m of noise lands them outside their own limits. # Under clamping that starts about half of all episodes with the jaws # less than fully open and about a tenth of them fully shut — in the # demo whose entire thesis is that a shut jaw cannot get AROUND a 5 cm # block. It also contradicts this file twice over: `arm_joint_pattern` # says the gripper owns the finger joints and nothing else may write to # them, and `on_start` says the controller and the gripper drive # DISJOINT joint sets. # # `self.gripper.reset(env_ids)` below does NOT undo it. That call # writes joint position TARGETS; this writes actual joint STATE, so the # fingers still START from the perturbed width and merely travel back # toward open over the following steps. joint_pos = self.robot.internals.default_joint_pos[env_ids].clone() arm = self._arm_dof_idx joint_pos[:, arm] += torch.empty_like(joint_pos[:, arm]).uniform_(-0.05, 0.05) joint_vel = self.robot.internals.default_joint_vel[env_ids] self.robot.set_joint_state(joint_pos, velocities=joint_vel, env_ids=env_ids)
# Re-randomise the block — the reason this task can exist. default_pose # is the spawn pose in the WORLD frame (env origins already applied), # so the jitter below is all that has to be added, and only the # environments in env_ids are touched: the ones still mid-episode keep # the block exactly where they left it. count = len(env_ids) pose = self.block.default_pose[env_ids].clone() pose[:, 0] += torch.empty(count, device=self.device).uniform_(-self.block_x_jitter, self.block_x_jitter) pose[:, 1] += torch.empty(count, device=self.device).uniform_(-self.block_y_jitter, self.block_y_jitter) self.block.set_pose(pose, env_ids=env_ids) # Zero the velocity too, or a block that was falling when the episode # ended arrives at its new pose still falling. self.block.set_velocity(torch.zeros(count, 6, device=self.device), env_ids=env_ids)
self.ik.reset()
# Open the gripper for exactly the environments being reset — passing # env_ids matters here: resetting every environment would unlatch the # thousands still mid-episode and drop whatever they were holding. self.gripper.reset(env_ids)
self._target_pos[env_ids] = self._initial_target[env_ids] self._actions[env_ids] = 0.0 # Shut the gate for the environments starting a fresh episode. Without # this a policy that ended one episode holding the block would begin # the next one already permitted to close, 28 cm away from a block that # has just been re-randomised somewhere else — which is the exact # behaviour the gate exists to remove. self._gate_open[env_ids] = False # ...and clear the refusal bit with it, so a fresh episode's first # reward cannot be charged for a close the PREVIOUS episode commanded. self._futile_close[env_ids] = False
# Put the goal above where the block was just placed — computed from # the pose that was just WRITTEN, never from a readback. `set_pose` # lands in the simulation's buffer and is not visible through # `state.pose` until the next physics step, so reading it back here # would sample the goal above the PREVIOUS episode's block. self._place_goals(env_ids, pose[:, :3] - self._base_frame_offset()[env_ids])
@app.job( # Tier 1: T4 GPU, 16 GB VRAM. Run `simulo systems` for the full four-tier catalog. system=simulo.SystemType.TIER_1, timeout=4 * 60 * 60, retries=2, callbacks=[simulo.callbacks.ResumableCheckpoint(every=50)],)def train_franka_lift(num_envs: int = 2048, max_iterations: int = 500) -> dict[str, Any]: """Train a pick-and-lift policy with PPO and save the checkpoint.
Args: num_envs: Number of parallel environments to simulate. max_iterations: Number of PPO policy-update iterations.
Returns: A JSON-serialisable dict: the saved ``checkpoint`` path plus training ``stats`` (``iterations``, ``num_envs``).
Raises: ValueError: If ``num_envs`` or ``max_iterations`` is not a positive integer. """ if num_envs < 1: raise ValueError(f"num_envs must be a positive integer, got {num_envs}") if max_iterations < 1: raise ValueError(f"max_iterations must be a positive integer, got {max_iterations}")
env = simulo.LearningEnv( task=FrankaLiftTask(), num_envs=num_envs, device="cuda", dt=1.0 / 120.0, physics_steps_per_action=2, env_spacing=2.5, headless=True, seed=42, ) trainer = simulo.RLTrainer( env=env, algorithm="PPO", device="cuda", seed=42, # PPO here defaults to NO exploration bonus (entropy_loss_scale 0.0), # which is a fine default for short-horizon tasks and a doubtful one # for manipulation, where a policy can converge onto a local optimum # and stop sampling the behaviour you want. # # Honest about the evidence: on THIS task, two 1500-iteration runs with # and without it did not separate on any metric measured — same final # reward to within run-to-run noise, same number of complete lifts # (two, each in a single environment). It is kept because it is the # right default for the task shape and costs nothing, NOT because it # was shown to help here. Do not read it as a fix. agent_cfg={"entropy_loss_scale": 0.01}, )
stats = trainer.train(max_iterations=max_iterations)
# TWO checkpoints, deliberately, because they answer different questions. # `franka_lift_final.pt` is the LAST policy — the one to resume from. # `franka_lift_best.pt` is the best-so-far checkpoint the trainer tracked; on # the measured 1500-iteration run it scored 98.00% against the last policy's # 92.97%, so it is the one to EVALUATE. Publishing only the last policy meant # whoever pulled the volume got the worse of the two with nothing saying so. # Read the best-vs-last discussion in this module's docstring before treating # `best.pt` as authoritative — on this task the tracking signal is coarse. checkpoint = f"{vol.path}/franka_lift_final.pt" trainer.save(checkpoint)
best_checkpoint = None checkpoint_dir = stats.get("checkpoint_dir") if checkpoint_dir: best_source = os.path.join(checkpoint_dir, "best.pt") # Absent whenever no chunk ever measured a finished episode — a normal # state, not an error, so this stays best-effort rather than asserting. if os.path.isfile(best_source): best_checkpoint = f"{vol.path}/franka_lift_best.pt" shutil.copyfile(best_source, best_checkpoint)
# Close the trainer before the environment so the RL library releases its # resources first. trainer.close() env.close()
return { "checkpoint": checkpoint, "best_checkpoint": best_checkpoint, "num_envs": num_envs, **stats, }Run it:
simulo run franka_lift/app.py --num-envs 2048 --max-iterations 500The completed result reports lift performance and names the saved checkpoint. The app demonstrates three manipulation patterns you can carry into other tasks:
Propresets the block per environment and exposes its live pose for rewards.ParallelGripperActuatorturns one policy action into open, close, or hold.- Object state defines success, while gripper state remains an observation.