Add robot_arm

build.sh

@@ -4,6 +4,7 @@ set -e
 # Usage:
 #   ./build.sh breakout              # Full native train/eval binary (CPU envs)
 #   ./build.sh breakout --gpu        # GPU env (ENV_HEADER=ocean/ENV/ENV.cu; exclusive vs .h)
+#   ./build.sh robot_arm             # CUDA-only; implies --gpu
 #   ./build.sh breakout --float      # float32 precision (required for --slowly)
 #   ./build.sh breakout --cpu        # Tiny standalone CPU eval executable
 #   ./build.sh breakout --debug      # Debug build
@@ -35,6 +36,16 @@ for arg in "$@"; do
     esac
 done
 
+if [ "$ENV" = "robot_arm" ]; then
+    USE_GPU_ENV=1
+    case "${MODE:-native}" in
+        local|fast|web|cpu)
+            echo "Error: robot_arm physics is CUDA-only; use the native GPU build" >&2
+            exit 1
+            ;;
+    esac
+fi
+
 if [ "$ENV" = "all" ]; then
     FAILED=""
     for env_dir in ocean/*/; do

config/robot_arm.ini

@@ -0,0 +1,40 @@
+[base]
+env_name = robot_arm
+cudagraphs = 1
+seed = 42
+checkpoint_interval = 10
+
+[vec]
+total_agents = 4096
+num_buffers = 1
+num_threads = 4
+action_mask_size = 0
+gpu_env = 1
+
+[env]
+model_glb = resources/robot_arm/franka_panda.glb
+no_timeout = 0
+stack = 0
+basketball = 0
+
+[policy]
+hidden_size = 256
+num_layers = 2
+expansion_factor = 2
+
+[train]
+total_timesteps = 1_000_000_000
+learning_rate = 0.0003
+anneal_lr = 1
+min_lr_ratio = 0.0
+gamma = 0.99
+gae_lambda = 0.95
+horizon = 128
+minibatch_size = 32768
+replay_ratio = 1.0
+clip_coef = 0.2
+ent_coef = 0.002
+anneal_ent_coef = 1
+min_ent_coef_ratio = 0.5
+vf_coef = 0.5
+max_grad_norm = 0.5

ocean/robot_arm/robot_arm.cu

@@ -0,0 +1,140 @@
+#ifndef PUFFER_ROBOT_ARM_GPU_CU
+#define PUFFER_ROBOT_ARM_GPU_CU
+
+#define PUF_BACKEND PUF_GPU
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+typedef __nv_bfloat16 obs_t;
+
+#include "robot_arm_cuda.cuh"
+
+static int g_ra_no_timeout;
+static int g_ra_stack;
+static int g_ra_basketball;
+static const char* g_ra_model_glb = "resources/robot_arm/franka_panda.glb";
+static RaRenderHost g_ra_render_host;
+
+static struct {
+    Env* envs;
+    int n;
+    obs_t* observations;
+    float* actions;
+    float* rewards;
+    float* terminals;
+    cudaStream_t stream;
+} g_gpu;
+
+static int ra_flag(Dict* kwargs, const char* key) {
+    DictItem* item = dict_find(kwargs, key);
+    return item != NULL && item->value != 0.0;
+}
+
+static void ra_fill(Env* env, unsigned int rng) {
+    memset(env, 0, sizeof(*env));
+    env->num_agents = 1;
+    env->rng = rng ? rng : 1u;
+    env->world.state.rng = env->rng;
+    env->world.state.no_timeout = g_ra_no_timeout;
+    env->world.state.stack_mode = g_ra_stack;
+    env->world.state.basketball_mode = g_ra_basketball;
+    ra_reset(&env->world.state);
+    ra_rbrst(&env->world.rigid, ra_topo(&env->world.state));
+}
+
+Env* puf_vec_create(int n, Dict* env_kwargs,
+        obs_t* observations, float* actions, float* rewards, float* terminals) {
+    g_ra_no_timeout = ra_flag(env_kwargs, "no_timeout");
+    g_ra_stack = ra_flag(env_kwargs, "stack");
+    g_ra_basketball = ra_flag(env_kwargs, "basketball");
+    assert(!(g_ra_stack && g_ra_basketball));
+    DictItem* model = dict_find(env_kwargs, "model_glb");
+    if (model != NULL && model->str != NULL && model->str[0] != '\0'
+            && strcmp(model->str, "None") != 0) {
+        g_ra_model_glb = model->str;
+    }
+    g_ra_render_host.model_glb = g_ra_model_glb;
+    g_ra_render_host.camera_distance = g_ra_basketball ? 2.35f : 1.55f;
+    g_ra_render_host.camera_yaw = 0.78f;
+    g_ra_render_host.camera_pitch = 0.48f;
+
+    Env* host_envs = (Env*)calloc(n, sizeof(Env));
+    for (int i = 0; i < n; i++) {
+        ra_fill(&host_envs[i], i + 1);
+    }
+    Env* envs = NULL;
+    assert(cudaMalloc((void**)&envs, n * sizeof(Env)) == cudaSuccess);
+    assert(cudaMemcpy(envs, host_envs, n * sizeof(Env),
+        cudaMemcpyHostToDevice) == cudaSuccess);
+    free(host_envs);
+    g_gpu.envs = envs;
+    g_gpu.n = n;
+    g_gpu.observations = observations;
+    g_gpu.actions = actions;
+    g_gpu.rewards = rewards;
+    g_gpu.terminals = terminals;
+    g_gpu.stream = 0;
+    return envs;
+}
+
+void puf_bind_stream(cudaStream_t stream) {
+    g_gpu.stream = stream;
+}
+
+void puf_init(Env*, Dict*) {
+}
+
+void puf_reset(Env*) {
+    ra_kinit<<<(g_gpu.n + RA_CUDA_BLOCK_SIZE - 1) / RA_CUDA_BLOCK_SIZE,
+        RA_CUDA_BLOCK_SIZE>>>(
+        g_gpu.envs, g_gpu.observations, g_gpu.rewards, g_gpu.terminals,
+        g_gpu.n);
+    assert(cudaGetLastError() == cudaSuccess);
+}
+
+void puf_step(Env*) {
+    dim3 grid((g_gpu.n + RA_CUDA_BLOCK_SIZE - 1) / RA_CUDA_BLOCK_SIZE);
+    dim3 block(RA_CUDA_BLOCK_SIZE);
+    ra_kbegin<<<grid, block, 0, g_gpu.stream>>>(
+        g_gpu.envs, 0, g_gpu.n, g_gpu.actions);
+    assert(cudaGetLastError() == cudaSuccess);
+    ra_kphys<<<grid, block, 0, g_gpu.stream>>>(
+        g_gpu.envs, 0, g_gpu.n);
+    assert(cudaGetLastError() == cudaSuccess);
+    ra_kfin<<<grid, block, 0, g_gpu.stream>>>(
+        g_gpu.envs, 0, g_gpu.n, g_gpu.observations, g_gpu.rewards,
+        g_gpu.terminals);
+    assert(cudaGetLastError() == cudaSuccess);
+}
+
+void puf_close(Env*) {
+    ra_rclose(&g_ra_render_host);
+    cudaFree(g_gpu.envs);
+    g_gpu.envs = NULL;
+}
+
+void puf_render(Env*) {
+    if (g_gpu.stream) {
+        cudaStreamSynchronize(g_gpu.stream);
+    }
+    RaState state;
+    assert(cudaMemcpy(&state, &g_gpu.envs->world.state, sizeof(RaState),
+        cudaMemcpyDeviceToHost) == cudaSuccess);
+    RaPose links[RA_LINKS];
+    ra_fk(state.q, state.gripper_width, links, NULL, NULL, &state.end_effector);
+    ra_draw(&g_ra_render_host, &state, links);
+    if (!g_ra_render_host.reset_requested) {
+        return;
+    }
+    g_ra_render_host.reset_requested = 0;
+    Env host_env;
+    ra_fill(&host_env, state.rng ? state.rng : 1u);
+    assert(cudaMemcpy(g_gpu.envs, &host_env, sizeof(Env),
+        cudaMemcpyHostToDevice) == cudaSuccess);
+}
+
+#endif