Building the RL Pipeline

Every call to env.step(actions) runs a fixed pipeline: apply the actions, step the physics, compute observations, compute rewards, check termination, and reset any finished environments. In this lesson we’ll walk through each stage of the Franka cube-lift task so you understand exactly what the agent senses, does, and is rewarded for.

In this lesson, we will:

• Apply policy actions with per-joint scaling and a staged gripper.

• Describe the reach-grasp-lift reward and the terms that shape it.

• Reset environments with randomized arm and cube positions.

• Assemble the 41-dimensional observation vector.

Apply Actions

_pre_physics_step clamps raw policy outputs to [-1, 1], applies per-joint arm scales relative to the default pose, and keeps the gripper open until the cube is centered well enough between the fingertips to start closing. _apply_action then stabilizes the robot state and writes the joint position targets.

notebook_env_cfg.action_scale = (0.45, 1.60, 0.70, 2.70, 0.45, 0.80, 0.30)


def _pre_physics_step(self, actions: torch.Tensor) -> None:
    self.actions = actions.clone().clamp(-1.0, 1.0)

    arm_targets = self.robot_default_joint_pos[:, self.arm_joint_indices] + self.arm_action_scale.unsqueeze(
        0
    ) * self.actions[:, : len(self.arm_joint_indices)]
    self.robot_dof_targets[:, self.arm_joint_indices] = torch.clamp(
        arm_targets,
        self.arm_dof_lower_limits.unsqueeze(0),
        self.arm_dof_upper_limits.unsqueeze(0),
    )

    allow_close_control = self.gripper_close_enabled
    self.actions[:, 7] = torch.where(allow_close_control, self.actions[:, 7], torch.ones_like(self.actions[:, 7]))

    finger_cmd = 0.5 * (self.actions[:, 7] + 1.0)
    finger_target = self.gripper_close_pos + finger_cmd * (self.gripper_open_pos - self.gripper_close_pos)
    self.robot_dof_targets[:, self.finger_joint_indices] = finger_target.unsqueeze(-1).expand(
        -1, len(self.finger_joint_indices)
    )


def _apply_action(self) -> None:
    self._stabilize_robot_state()
    self.robot.set_joint_position_target_index(target=self.robot_dof_targets)

Gating the gripper (gripper_close_enabled) is a small but important detail: it stops the policy from closing the fingers before the cube is well positioned, which prevents the agent from learning to bat the cube away.

Rewards

_get_rewards implements a reach-grasp-lift pipeline: reach the cube, center it between the fingertips while the gripper stays open, close into a grasp, and then lift. The reward is a weighted sum of stage terms plus small penalties on action magnitude and joint velocity.

notebook_env_cfg.reaching_object_scale = 1.0
notebook_env_cfg.reaching_object_std = 0.1
notebook_env_cfg.gripper_open_reward_scale = 0.2
notebook_env_cfg.grasp_reward_scale = 4.0
notebook_env_cfg.lifting_object_scale = 16.0
notebook_env_cfg.lift_progress_reward_scale = 8.0
notebook_env_cfg.action_penalty_scale = 1e-4
notebook_env_cfg.joint_vel_penalty_scale = 1e-4
notebook_env_cfg.lifted_height = 0.08

The environment computes grasp-quality metrics (how centered, balanced, and enclosed the cube is between the fingers) and uses them to gate the later stages. Once the cube is reached and well positioned, gripper_close_enabled is turned on; once grasped, incremental lifting is rewarded. The stage terms are combined like this:

rewards = (
    self.cfg.reaching_object_scale * reaching_reward
    + self.cfg.gripper_open_reward_scale * enclosure_reward
    + self.cfg.grasp_reward_scale * grasp_reward
    + self.cfg.lift_progress_reward_scale * lift_progress_reward
    + self.cfg.lifting_object_scale * lifting_reward
    - self.cfg.action_penalty_scale * action_penalty
    - self.cfg.joint_vel_penalty_scale * joint_vel_penalty
)
rewards = torch.nan_to_num(rewards, nan=0.0, posinf=0.0, neginf=0.0)

Tip

The relative weights encode the task’s priorities: lifting (16.0) dominates once a grasp is secured, while reaching (1.0) is a gentle, always-on shaping signal that guides the arm early in training. The full grasp-metric helpers live in the franka_cube extension source.

Reset

_reset_idx resets the robot and cube state, adds small arm and cube position noise for variety, and reopens the gripper. Randomizing initial conditions is what forces the policy to generalize instead of memorizing one trajectory.

notebook_env_cfg.episode_length_s = 3.0
notebook_env_cfg.object_drop_height = -0.05
notebook_env_cfg.object_reset_pos_x_range = (-0.08, 0.03)
notebook_env_cfg.object_reset_pos_y_range = (-0.2, 0.2)
notebook_env_cfg.reset_arm_noise = 0.1

On reset, the environment restores the default root pose and velocity, applies uniform noise to the arm joints (clamped to limits), reopens the fingers, and re-randomizes the cube’s x and y position within the configured ranges. It also clears the action buffers and resets gripper_close_enabled to False so the staged gripper logic starts fresh each episode.

Observations

_get_observations builds a 41-dimensional observation vector per environment:

Slice	Dim	Description
[0:9]	9	Joint positions relative to default pose
[9:18]	9	Joint velocities
[18:21]	3	Cube position relative to the grasp frame
[21:24]	3	Cube position relative to the left fingertip
[24:27]	3	Cube position relative to the right fingertip
[27:30]	3	Cube position relative to the fingertip midpoint
[30:38]	8	Previous actions
[38:39]	1	Finger open fraction
[39:40]	1	Enclosure gate
[40:41]	1	Grasp reward proxy

def _get_observations(self) -> dict:
    self._stabilize_robot_state()

    joint_pos = self.robot.data.joint_pos.torch
    joint_vel = self.robot.data.joint_vel.torch
    object_pos = self._get_object_pos()
    grasp_pos = self._get_grasp_pos()
    left_finger_pos, right_finger_pos = self._get_finger_positions()
    finger_joint_pos = joint_pos[:, self.finger_joint_ids].mean(dim=-1)
    finger_midpoint, _, enclosure_gate, grasp_reward, finger_open_fraction = self._compute_grasp_metrics(
        object_pos=object_pos,
        left_finger_pos=left_finger_pos,
        right_finger_pos=right_finger_pos,
        finger_joint_pos=finger_joint_pos,
    )

    obs = torch.cat(
        [
            joint_pos - self.robot_default_joint_pos,
            joint_vel,
            object_pos - grasp_pos,
            object_pos - left_finger_pos,
            object_pos - right_finger_pos,
            object_pos - finger_midpoint,
            self.previous_actions,
            finger_open_fraction.unsqueeze(-1),
            enclosure_gate.unsqueeze(-1),
            grasp_reward.unsqueeze(-1),
        ],
        dim=-1,
    )
    obs = torch.nan_to_num(obs, nan=0.0, posinf=100.0, neginf=-100.0)
    obs = torch.clamp(obs, -100.0, 100.0)
    self.previous_actions[:] = self.actions
    return {"policy": obs}

Notice the observation is expressed in relative terms (positions relative to the default pose and to the fingertips) rather than absolute world coordinates. Relative features generalize better and make the policy robust to where the cube spawns.

Key Takeaways

You now understand every stage of the RL pipeline for the Franka cube-lift task: how actions are scaled and gated, how the reach-grasp-lift reward is composed, how resets randomize the scene, and how the 41-dimensional relative observation is assembled. These design choices are the difference between a policy that learns and one that flails. Next, we’ll run the full loop and then load a trained policy to see the payoff.