CPU fallback for mac scrubs

Author: jsuarez5341

build.sh

@@ -10,7 +10,7 @@ set -e
 #   ./build.sh breakout --web        # Emscripten web build
 #   ./build.sh breakout --profile    # Kernel profiling binary
 
-ENV=${1:?Usage: ./build.sh ENV_NAME [--float] [--debug] [--local|--fast|--web|--profile]}
+ENV=${1:?Usage: ./build.sh ENV_NAME [--float] [--debug] [--local|--fast|--web|--profile|--cpu]}
 MODE=""
 PRECISION=""
 DEBUG=""
@@ -22,6 +22,7 @@ for arg in "${@:2}"; do
         --fast)  MODE=fast ;;
         --web)   MODE=web ;;
         --profile) MODE=profile ;;
+        --cpu)   MODE=cpu; PRECISION="-DPRECISION_FLOAT" ;;
     esac
 done
 
@@ -203,6 +204,33 @@ if [ "$MODE" = "profile" ]; then
     exit 0
 fi
 
+if [ "$MODE" = "cpu" ]; then
+    echo "=== Compiling bindings_cpu.cpp ==="
+    g++ -c -fPIC -fopenmp \
+        -D_GLIBCXX_USE_CXX11_ABI=1 \
+        -DPLATFORM_DESKTOP \
+        -std=c++17 \
+        -I. -Isrc \
+        -I$PYTHON_INCLUDE -I$PYBIND_INCLUDE \
+        -DOBS_TENSOR_T=$OBS_TENSOR_T \
+        $PRECISION $LINK_OPT \
+        src/bindings_cpu.cpp -o src/bindings_cpu.o
+
+    echo "=== Linking $OUTPUT (CPU) ==="
+    LINK_CMD=(
+        g++ -shared -fPIC -fopenmp
+        src/bindings_cpu.o "$STATIC_LIB" "$RAYLIB_A"
+        -lm -lpthread -lomp5
+        $LINK_OPT
+    )
+    [ "$PLATFORM" = "Linux" ] && LINK_CMD+=(-Bsymbolic-functions)
+    [ "$PLATFORM" = "Darwin" ] && LINK_CMD+=(-framework Cocoa -framework OpenGL -framework IOKit)
+    LINK_CMD+=(-o "$OUTPUT")
+    "${LINK_CMD[@]}"
+    echo "=== Built: $OUTPUT (CPU) ==="
+    exit 0
+fi
+
 echo "=== Compiling bindings.cu ==="
 $NVCC -c -Xcompiler -fPIC \
     -Xcompiler=-D_GLIBCXX_USE_CXX11_ABI=1 \

constellation/cache_data.py

@@ -188,23 +188,11 @@ def cached_load(path, env_name, cache):
     #data['metrics/agent_steps'] = [e/1e6 for e in data['metrics/agent_steps']]
     del data['metrics/agent_steps']
 
-    '''
-    for k, v in data.items():
-        for e in v:
-            if e is None or isinstance(e, str):
-                continue
-            try:
-                if e > 1e9 or e < -1e9:
-                    breakpoint()
-            except:
-                breakpoint()
-    '''
-
     # Filter to pareto
+    '''
     steps = data['agent_steps']
     costs = data['uptime']
     scores = data['env/score']
-    '''
     idxs = pareto_idx(steps, costs, scores)
     for k in data:
         try:

pufferlib/torch_pufferl.py

@@ -5,6 +5,7 @@
 import os
 import glob
 import time
+import ctypes
 from collections import defaultdict
 
 import numpy as np
@@ -93,6 +94,20 @@ def __init__(self, ptr, shape, dtype):
             'version': 2,
         }
 
+_TORCH_TO_CTYPE = {
+    torch.uint8:   ctypes.c_uint8,
+    torch.float32: ctypes.c_float,
+}
+
+def _cpu_tensor(ptr, shape, dtype):
+    '''Zero-copy CPU tensor from a raw pointer via ctypes.'''
+    ctype = _TORCH_TO_CTYPE[dtype]
+    n = 1
+    for s in shape:
+        n *= s
+    arr = (ctype * n).from_address(ptr)
+    return torch.frombuffer(arr, dtype=dtype).reshape(shape)
+
 class PuffeRL:
     def __init__(self, args, vec, policy, verbose=True):
         config = args['train']
@@ -103,16 +118,25 @@ def __init__(self, args, vec, policy, verbose=True):
         torch.backends.cudnn.benchmark = True
 
         self._vec = vec
+        self.gpu = vec.gpu
         total_agents = vec.total_agents
         self.total_agents = total_agents
         obs_dtype = _OBS_DTYPE_MAP.get(vec.obs_dtype, torch.uint8)
 
-        self.vec_obs = torch.as_tensor(_CudaPtr(vec.gpu_obs_ptr,
-            (total_agents, vec.obs_size), obs_dtype))
-        self.vec_rewards = torch.as_tensor(_CudaPtr(vec.gpu_rewards_ptr,
-            (total_agents,), torch.float32))
-        self.vec_terminals = torch.as_tensor(_CudaPtr(vec.gpu_terminals_ptr,
-            (total_agents,), torch.float32))
+        if self.gpu:
+            self.vec_obs = torch.as_tensor(_CudaPtr(vec.gpu_obs_ptr,
+                (total_agents, vec.obs_size), obs_dtype))
+            self.vec_rewards = torch.as_tensor(_CudaPtr(vec.gpu_rewards_ptr,
+                (total_agents,), torch.float32))
+            self.vec_terminals = torch.as_tensor(_CudaPtr(vec.gpu_terminals_ptr,
+                (total_agents,), torch.float32))
+        else:
+            self.vec_obs = _cpu_tensor(vec.obs_ptr,
+                (total_agents, vec.obs_size), obs_dtype)
+            self.vec_rewards = _cpu_tensor(vec.rewards_ptr,
+                (total_agents,), torch.float32)
+            self.vec_terminals = _cpu_tensor(vec.terminals_ptr,
+                (total_agents,), torch.float32)
 
         vec.reset()
         horizon = config['horizon']
@@ -148,7 +172,7 @@ def __init__(self, args, vec, policy, verbose=True):
         self.last_log_step = 0
         self.last_log_time = time.time()
         self.start_time = time.time()
-        self.profile = Profile()
+        self.profile = Profile(gpu=self.gpu)
         self.verbose = verbose
 
         self.model_size = sum(p.numel() for p in policy.parameters() if p.requires_grad)
@@ -203,9 +227,13 @@ def rollouts(self):
                 self.values[t] = value.flatten()
 
             prof.mark(2)
-            actions_gpu = (action.T if action.dim() > 1 else action.unsqueeze(-1)).to(dtype=torch.float32, device='cuda').contiguous()
-            self._vec.step(actions_gpu.data_ptr())
-            torch.cuda.synchronize()
+            actions_flat = (action.T if action.dim() > 1 else action.unsqueeze(-1)).to(dtype=torch.float32).contiguous()
+            if self.gpu:
+                actions_flat = actions_flat.cuda()
+                self._vec.gpu_step(actions_flat.data_ptr())
+                torch.cuda.synchronize()
+            else:
+                self._vec.cpu_step(actions_flat.data_ptr())
             o, r, d = self.vec_obs, self.vec_rewards, self.vec_terminals
             prof.mark(3)
 
@@ -348,7 +376,7 @@ def log(self):
                 'train_misc': perf[P.TRAIN_MISC],
                 'train_forward': perf[P.TRAIN_FORWARD],
             },
-            'util': dict(_C.get_utilization(self.args.get('gpu_id', 0))),
+            'util': dict(_C.get_utilization(self.args.get('gpu_id', 0))) if self.gpu else {},
         }
         self.last_log_time = time.time()
         self.last_log_step = self.global_step
@@ -376,7 +404,8 @@ def create_pufferl(cls, args):
             os.environ['CUDA_VISIBLE_DEVICES'] = str(local_rank)
 
         args['vec']['num_buffers'] = 1
-        vec = _C.create_vec(args)
+        gpu = 1 if device == 'cuda' else 0
+        vec = _C.create_vec(args, gpu)
         policy = load_policy(args, vec)
 
         if 'LOCAL_RANK' in os.environ:
@@ -395,7 +424,8 @@ def create_pufferl(cls, args):
 def compute_puff_advantage(values, rewards, terminals,
         ratio, advantages, gamma, gae_lambda, vtrace_rho_clip, vtrace_c_clip):
     num_steps, horizon = values.shape
-    _C.puff_advantage(
+    fn = _C.puff_advantage if values.is_cuda else _C.puff_advantage_cpu
+    fn(
         values.data_ptr(), rewards.data_ptr(), terminals.data_ptr(),
         ratio.data_ptr(), advantages.data_ptr(),
         num_steps, horizon,
@@ -406,16 +436,26 @@ class Profile:
     '''Matches pufferlib.cu profiling: accumulate ms, report seconds.'''
     ROLLOUT, EVAL_GPU, EVAL_ENV, TRAIN, TRAIN_MISC, TRAIN_FORWARD, NUM = range(7)
 
-    def __init__(self):
+    def __init__(self, gpu=True):
         self.accum = [0.0] * Profile.NUM
-        self._events = [torch.cuda.Event(enable_timing=True) for _ in range(4)]
+        self.gpu = gpu
+        if gpu:
+            self._events = [torch.cuda.Event(enable_timing=True) for _ in range(4)]
+        else:
+            self._stamps = [0.0] * 4
 
     def mark(self, idx):
-        self._events[idx].record()
+        if self.gpu:
+            self._events[idx].record()
+        else:
+            self._stamps[idx] = time.perf_counter()
 
     def elapsed(self, idx, start_ev, end_ev):
-        self._events[end_ev].synchronize()
-        self.accum[idx] += self._events[start_ev].elapsed_time(self._events[end_ev])
+        if self.gpu:
+            self._events[end_ev].synchronize()
+            self.accum[idx] += self._events[start_ev].elapsed_time(self._events[end_ev])
+        else:
+            self.accum[idx] += (self._stamps[end_ev] - self._stamps[start_ev]) * 1000.0
 
     def read_and_reset(self):
         out = [v / 1000.0 for v in self.accum]

src/bindings.cu

@@ -256,9 +256,10 @@ struct VecEnv {
     std::vector<int> act_sizes;
     std::string obs_dtype;
     size_t obs_elem_size;
+    int gpu;
 };
 
-std::unique_ptr<VecEnv> create_vec(py::dict args) {
+std::unique_ptr<VecEnv> create_vec(py::dict args, int gpu) {
     py::dict vec_kwargs = args["vec"].cast<py::dict>();
     py::dict env_kwargs = args["env"].cast<py::dict>();
 
@@ -269,9 +270,10 @@ std::unique_ptr<VecEnv> create_vec(py::dict args) {
     Dict* env_dict = py_dict_to_c_dict(env_kwargs);
 
     auto ve = std::make_unique<VecEnv>();
+    ve->gpu = gpu;
     {
         py::gil_scoped_release no_gil;
-        ve->vec = create_static_vec(total_agents, num_buffers, vec_dict, env_dict);
+        ve->vec = create_static_vec(total_agents, num_buffers, gpu, vec_dict, env_dict);
     }
     ve->total_agents  = total_agents;
     ve->obs_size      = get_obs_size();
@@ -291,14 +293,22 @@ void vec_reset(VecEnv& ve) {
     static_vec_reset(ve.vec);
 }
 
-// actions_ptr: data_ptr() of a (total_agents, num_atns) float64 CUDA tensor
-void vec_step(VecEnv& ve, long long actions_ptr) {
+void gpu_vec_step_py(VecEnv& ve, long long actions_ptr) {
     cudaMemcpy(ve.vec->gpu_actions, (void*)actions_ptr,
         (size_t)ve.total_agents * ve.num_atns * sizeof(float),
         cudaMemcpyDeviceToDevice);
     {
         py::gil_scoped_release no_gil;
-        static_vec_step(ve.vec);
+        gpu_vec_step(ve.vec);
+    }
+}
+
+void cpu_vec_step_py(VecEnv& ve, long long actions_ptr) {
+    memcpy(ve.vec->actions, (void*)actions_ptr,
+        (size_t)ve.total_agents * ve.num_atns * sizeof(float));
+    {
+        py::gil_scoped_release no_gil;
+        cpu_vec_step(ve.vec);
     }
 }
 
@@ -512,20 +522,26 @@ PYBIND11_MODULE(_C, m) {
         return now - pufferl.start_time;
     });
     m.def("puff_advantage", &py_puff_advantage);
-    m.def("create_vec", &create_vec);
+    m.def("create_vec", &create_vec, py::arg("args"), py::arg("gpu") = 1);
     py::class_<VecEnv, std::unique_ptr<VecEnv>>(m, "VecEnv")
         .def_readonly("total_agents",  &VecEnv::total_agents)
         .def_readonly("obs_size",      &VecEnv::obs_size)
         .def_readonly("num_atns",      &VecEnv::num_atns)
         .def_readonly("act_sizes",     &VecEnv::act_sizes)
         .def_readonly("obs_dtype",     &VecEnv::obs_dtype)
         .def_readonly("obs_elem_size", &VecEnv::obs_elem_size)
+        .def_readonly("gpu",           &VecEnv::gpu)
         // GPU buffer pointers — wrap with torch.from_blob(..., device='cuda')
         .def_property_readonly("gpu_obs_ptr",       [](VecEnv& ve) { return (long long)ve.vec->gpu_observations; })
         .def_property_readonly("gpu_rewards_ptr",   [](VecEnv& ve) { return (long long)ve.vec->gpu_rewards; })
         .def_property_readonly("gpu_terminals_ptr", [](VecEnv& ve) { return (long long)ve.vec->gpu_terminals; })
+        // CPU buffer pointers (same as gpu_ in CPU mode since they alias)
+        .def_property_readonly("obs_ptr",       [](VecEnv& ve) { return (long long)ve.vec->observations; })
+        .def_property_readonly("rewards_ptr",   [](VecEnv& ve) { return (long long)ve.vec->rewards; })
+        .def_property_readonly("terminals_ptr", [](VecEnv& ve) { return (long long)ve.vec->terminals; })
         .def("reset", &vec_reset)
-        .def("step",  &vec_step)
+        .def("gpu_step", &gpu_vec_step_py)
+        .def("cpu_step", &cpu_vec_step_py)
         .def("log",   &vec_log)
         .def("close", &vec_close);