apiVersion: v1
kind: Pod
metadata:
  generateName: nri-train-sparse-k6-
  labels:
    type: interactive
    research-group: rl-team
spec:
  hostIPC: true
  restartPolicy: Never
  containers:
  - name: nri-train
    command: ["python", "-u", "train.py", 
    "--epochs=50",
    "--suffix=_dance", 
    "--num-atoms=53",
    "--prediction-steps=10",
    "--dims=6",
    "--edge-types=6",
    "--encoder=mlp",
    "--decoder=rnn",
    "--dynamic-graph",
    "--skip-first",
    "--prior", # sparsity prior
#     "--hard",
    "--save-folder=/workspace/NRI/logs/dance_53joints_rnndecoder_sparse_k6",
  ]
    # command: ["sleep", "infinity"]
    image: amr-registry.caas.intel.com/aipg/mpettee-nri
    workingDir: /workspace/NRI/
    resources:
      limits:
        nvidia.com/gpu: 0
        cpu: 1
        memory: 350Gi
    volumeMounts:
    - mountPath: "/workspace"
      name: workspace 
    envFrom:
    - configMapRef:
        name: proxy-config
    env:
    - name: LD_LIBRARY_PATH
      value: $LD_LIBRARY_PATH:/.mujoco/mjpro150/bin
    - name: OMP_NUM_THREADS
      value: "1"
  securityContext:
    fsGroup: 17685
    runAsUser: 11964579
  volumes:
  - name: workspace
    persistentVolumeClaim:
      claimName: workdisk

